test: Tests for workflow async jobs API

writer · Dec 10, 2024 · fa0d470 · fa0d470
1 parent 278f668
commit fa0d470
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 25 deletions.
diff --git a/src/writer/crypto.py b/src/writer/crypto.py
@@ -6,7 +6,8 @@
 
 HASH_SALT = "a9zHYfIeL0"
 
-def get_hash(message: str, base_hash = os.getenv("WRITER_BASE_HASH")):
+def get_hash(message: str):
+    base_hash = os.getenv("WRITER_BASE_HASH")
     if not base_hash:
         raise ValueError("Environment variable WRITER_BASE_HASH needs to be set up in" + \
                             "order to enable operations which require hash generation, such as creating async jobs.")

diff --git a/src/writer/serve.py b/src/writer/serve.py
@@ -325,7 +325,7 @@ def update_job(job_id: str, job_info: dict):
             merged_info = current_job_info | { "finished_at": int(time.time()) } | job_info
             app.state.job_vault.set(job_id, merged_info)
 
-        def job_done_callback(task, job_id: str):
+        def job_done_callback(task: asyncio.Task, job_id: str):
             try:
                 apsr: Optional[AppProcessServerResponse] = None
                 apsr = task.result()

diff --git a/tests/backend/test_serve.py b/tests/backend/test_serve.py
@@ -1,5 +1,5 @@
 import mimetypes
-import os
+import time
 
 import fastapi
 import fastapi.testclient
@@ -215,28 +215,64 @@ def test_feature_flags(self):
                 "Content-Type": "application/json"
             })
             feature_flags = res.json().get("featureFlags")
-            assert feature_flags == ["flag_one", "flag_two"]
+            assert feature_flags == ["workflows", "flag_one", "flag_two"]
 
-    # def test_create_workflow_job_api(self):
-    #     asgi_app: fastapi.FastAPI = writer.serve.get_asgi_app(
-    #         test_app_dir, "run")
-    #     os.environ["WRITER_BASE_HASH"] = "abc"
-    #     workflow_key = "workflow2"
+    def test_create_workflow_job_api(self, monkeypatch):
+        asgi_app: fastapi.FastAPI = writer.serve.get_asgi_app(
+            test_app_dir, "run")
+        monkeypatch.setenv("WRITER_BASE_HASH", "abc")
+        workflow_key = "workflow2"
 
-    #     with fastapi.testclient.TestClient(asgi_app) as client:
-    #         create_job_token = crypto.get_hash(f"create_job_{workflow_key}")
-    #         res = client.post(f"/api/job/workflow/{workflow_key}", json={
-    #             "proposedSessionId": None
-    #         }, headers={
-    #             "Content-Type": "application/json",
-    #             "Authorization": f"Bearer {create_job_token}"
-    #         })
-    #         job_id = res.json().get("id")
-    #         get_job_token = res.json().get("token")
-    #         res = client.get(f"/api/job/{job_id}", headers={
-    #             "Authorization": f"Bearer {get_job_token}"
-    #         })
-    #         assert res.json().get("result") == 987127
-
-    #     os.environ["WRITER_BASE_HASH"] = ""
+        with fastapi.testclient.TestClient(asgi_app) as client:
+            create_job_token = crypto.get_hash(f"create_job_{workflow_key}")
+            res = client.post(f"/api/job/workflow/{workflow_key}", json={
+                "proposedSessionId": None
+            }, headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {create_job_token}"
+            })
+            time.sleep(1)
+            job_id = res.json().get("id")
+            get_job_token = res.json().get("token")
+            res = client.get(f"/api/job/{job_id}", headers={
+                "Authorization": f"Bearer {get_job_token}"
+            })
+            assert res.json().get("result") == "987127"
 
+    def test_create_workflow_job_api_incorrect_token(self, monkeypatch):
+        asgi_app: fastapi.FastAPI = writer.serve.get_asgi_app(
+            test_app_dir, "run")
+        monkeypatch.setenv("WRITER_BASE_HASH", "abc")
+        workflow_key = "workflow2"
+
+        with fastapi.testclient.TestClient(asgi_app) as client:
+            create_job_token = crypto.get_hash(f"not_the_right_message")
+            res = client.post(f"/api/job/workflow/{workflow_key}", json={
+                "proposedSessionId": None
+            }, headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {create_job_token}"
+            })
+            assert res.status_code == 403
+
+    def test_create_workflow_job_api_incorrect_token_for_get(self, monkeypatch):
+        asgi_app: fastapi.FastAPI = writer.serve.get_asgi_app(
+            test_app_dir, "run")
+        monkeypatch.setenv("WRITER_BASE_HASH", "abc")
+        workflow_key = "workflow2"
+
+        with fastapi.testclient.TestClient(asgi_app) as client:
+            create_job_token = crypto.get_hash(f"not_the_right_message")
+            res = client.post(f"/api/job/workflow/{workflow_key}", json={
+                "proposedSessionId": None
+            }, headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {create_job_token}"
+            })
+            time.sleep(1)
+            job_id = res.json().get("id")
+            get_job_token = "not_the_right_job_token"
+            res = client.get(f"/api/job/{job_id}", headers={
+                "Authorization": f"Bearer {get_job_token}"
+            })
+            assert res.status_code == 403