From fa0d4708377cad23b3b49ab6114444e4bf1791d6 Mon Sep 17 00:00:00 2001
From: Ramiro Medina <64783088+ramedina86@users.noreply.github.com>
Date: Tue, 10 Dec 2024 13:50:31 +0000
Subject: [PATCH] test: Tests for workflow async jobs API

---
 src/writer/crypto.py        |  3 +-
 src/writer/serve.py         |  2 +-
 tests/backend/test_serve.py | 82 ++++++++++++++++++++++++++-----------
 3 files changed, 62 insertions(+), 25 deletions(-)

diff --git a/src/writer/crypto.py b/src/writer/crypto.py
index 25be0f447..9333a264e 100644
--- 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
index 792cd3e5a..e412ee408 100644
--- 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
index d43341716..a39af0688 100644
--- 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
\ No newline at end of file