From 562a89753ed309540b4bd2c41d6c9ac727bd8f0f Mon Sep 17 00:00:00 2001
From: jcschaff <schaff.jim@gmail.com>
Date: Mon, 13 Jan 2025 16:53:15 -0500
Subject: [PATCH] wait until workflow complete in rest endpoint unit test

---
 biosim_server/api/main.py                          | 14 +++++++-------
 .../verify/workflows/omex_verify_workflow.py       |  2 +-
 tests/api/test_main.py                             | 14 ++++++++++++--
 3 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/biosim_server/api/main.py b/biosim_server/api/main.py
index 58043f2..2f5f022 100644
--- a/biosim_server/api/main.py
+++ b/biosim_server/api/main.py
@@ -115,14 +115,14 @@ def root() -> dict[str, str]:
 
 
 @app.post(
-    "/verify",
+    "/verify_omex",
     response_model=OmexVerifyWorkflowOutput,
     name="Uniform Time Course Comparison from OMEX/COMBINE archive",
-    operation_id="verify",
+    operation_id="start-verify-omex",
     tags=["Verification"],
     dependencies=[Depends(get_temporal_client), Depends(get_file_service)],
     summary="Compare UTC outputs from a deterministic SBML model within an OMEX/COMBINE archive.")
-async def verify(
+async def start_verify_omex(
         uploaded_file: UploadFile = File(..., description="OMEX/COMBINE archive containing a deterministic SBML model"),
         workflow_id_prefix: str = Query(default="verification-", description="Prefix for the workflow id."),
         simulators: list[str] = Query(default=["amici", "copasi", "pysces", "tellurium", "vcell"],
@@ -190,14 +190,14 @@ async def verify(
 
 
 @app.get(
-    "/get-output/{workflow_id}",
+    "/verify_omex/{workflow_id}",
     response_model=OmexVerifyWorkflowOutput,
-    operation_id='get-output',
+    operation_id='get-verify-omex',
     tags=["Results"],
     dependencies=[Depends(get_biosim_service), Depends(get_file_service)],
     summary='Get the results of an existing verification run.')
-async def get_output(workflow_id: str) -> OmexVerifyWorkflowOutput:
-    logger.info(f"in get /get-output/{workflow_id}")
+async def get_verify_omex(workflow_id: str) -> OmexVerifyWorkflowOutput:
+    logger.info(f"in get /verify_omex/{workflow_id}")
 
     try:
         # query temporal for the workflow output
diff --git a/biosim_server/verify/workflows/omex_verify_workflow.py b/biosim_server/verify/workflows/omex_verify_workflow.py
index ed83adf..5b2ad2f 100644
--- a/biosim_server/verify/workflows/omex_verify_workflow.py
+++ b/biosim_server/verify/workflows/omex_verify_workflow.py
@@ -123,5 +123,5 @@ async def run(self, verify_input: OmexVerifyWorkflowInput) -> OmexVerifyWorkflow
         )
         workflow.logger.info(f"Report generated at: {report_location}")
 
-        # Return the report location
+        self.verify_output.workflow_status = OmexVerifyWorkflowStatus.COMPLETED
         return self.verify_output
diff --git a/tests/api/test_main.py b/tests/api/test_main.py
index 6c53576..31a4847 100644
--- a/tests/api/test_main.py
+++ b/tests/api/test_main.py
@@ -1,3 +1,5 @@
+import asyncio
+import logging
 import os
 from copy import copy
 from pathlib import Path
@@ -30,7 +32,7 @@ async def test_get_output_not_found(verify_workflow_input: OmexVerifyWorkflowInp
 
     async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as test_client:
         # test with non-existent verification_id
-        response = await test_client.get(f"/get-output/non-existent-id")
+        response = await test_client.get(f"/verify_omex/non-existent-id")
         assert response.status_code == 404
 
 
@@ -59,7 +61,7 @@ async def test_verify_and_get_output(verify_workflow_input: OmexVerifyWorkflowIn
         with open(file_path, "rb") as file:
             files = {"uploaded_file": (
                 "BIOMD0000000010_tellurium_Negative_feedback_and_ultrasen.omex", file, "application/zip")}
-            response = await test_client.post("/verify", files=files, params=query_params)
+            response = await test_client.post("/verify_omex", files=files, params=query_params)
             assert response.status_code == 200
 
             workflow_output = OmexVerifyWorkflowOutput.model_validate(response.json())
@@ -80,6 +82,13 @@ async def test_verify_and_get_output(verify_workflow_input: OmexVerifyWorkflowIn
                 workflow_run_id=workflow_output.workflow_run_id
             )
 
+            while output.workflow_status != OmexVerifyWorkflowStatus.COMPLETED:
+                await asyncio.sleep(5)
+                response = await test_client.get(f"/verify_omex/{output.workflow_input.workflow_id}")
+                if response.status_code == 200:
+                    output = OmexVerifyWorkflowOutput.model_validate(response.json())
+                    logging.info(f"polling, job status is: {output.workflow_status}")
+
             # set timestamp, job_id, and omex_s3_path before comparison (these are set on server)
             expected_verify_workflow_output = OmexVerifyWorkflowOutput(
                 workflow_input=copy(verify_workflow_input),
@@ -91,6 +100,7 @@ async def test_verify_and_get_output(verify_workflow_input: OmexVerifyWorkflowIn
             expected_verify_workflow_output.workflow_input.workflow_id = output.workflow_input.workflow_id
             expected_verify_workflow_output.timestamp = output.timestamp
             expected_verify_workflow_output.workflow_run_id = output.workflow_run_id
+            expected_verify_workflow_output.workflow_status = OmexVerifyWorkflowStatus.COMPLETED
             assert expected_verify_workflow_output == output
 
     # verify the omex_s3_path file to the original file_path