feat: Fixed APIadapter (multiprocessing)

Alan Christie · Alan Christie · commit 6553b58b6f06 · 2024-11-04T16:06:46.000+01:00
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,4 @@
 .coverage
 dist/
 **/__pycache__/
+**/*.pickle
diff --git a/tests/api_adapter.py b/tests/api_adapter.py
@@ -1,4 +1,6 @@
 import os
+from multiprocessing import Lock
+from pickle import Pickler, Unpickler
 from typing import Any, Dict, List, Optional
 
 import yaml
@@ -14,61 +16,114 @@
 assert _JOB_DEFINITIONS
 
 _INSTANCE_ID_FORMAT: str = "instance-00000000-0000-0000-0000-{id:012d}"
+_TASK_ID_FORMAT: str = "task-00000000-0000-0000-0000-{id:012d}"
 _WORKFLOW_DEFINITION_ID_FORMAT: str = "workflow-00000000-0000-0000-0000-{id:012d}"
 _RUNNING_WORKFLOW_ID_FORMAT: str = "r-workflow-00000000-0000-0000-0000-{id:012d}"
 _RUNNING_WORKFLOW_STEP_ID_FORMAT: str = (
     "r-workflow-step-00000000-0000-0000-0000-{id:012d}"
 )
 
+_WORKFLOW_PICKLE_FILE: str = "workflow.pickle"
+_RUNNING_WORKFLOW_PICKLE_FILE: str = "running-workflow.pickle"
+_RUNNING_WORKFLOW_STEP_PICKLE_FILE: str = "running-workflow-step.pickle"
+_INSTANCE_PICKLE_FILE: str = "instance.pickle"
+_TASK_PICKLE_FILE: str = "task.pickle"
+
 
 class UnitTestAPIAdapter(APIAdapter):
     """A minimal API adapter. It serves-up Job Definitions
     from the job-definitions/job-definitions.yaml file and provides basic
     (in-memory) storage for Workflow Definitions and related tables."""
 
+    mp_lock = Lock()
+
     def __init__(self):
         super().__init__()
-        # A map of workflow definitions, keyed by workflow definition ID.
-        self._workflow_definitions: Dict[str, Dict[str, Any]] = {}
-        self._running_workflow: Dict[str, Dict[str, Any]] = {}
-        self._running_workflow_steps: Dict[str, Dict[str, Any]] = {}
-        self._instances: Dict[str, Dict[str, Any]] = {}
-        self._tasks: Dict[str, Dict[str, Any]] = {}
+        # Safely initialise the pickle files
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_WORKFLOW_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump({})
+        with open(f"tests/{_RUNNING_WORKFLOW_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump({})
+        with open(f"tests/{_RUNNING_WORKFLOW_STEP_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump({})
+        with open(f"tests/{_INSTANCE_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump({})
+        with open(f"tests/{_TASK_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump({})
+        UnitTestAPIAdapter.mp_lock.release()
 
     def create_workflow(self, *, workflow_definition: Dict[str, Any]) -> str:
-        next_id: int = len(self._workflow_definitions) + 1
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_WORKFLOW_PICKLE_FILE}", "rb") as pickle_file:
+            workflow = Unpickler(pickle_file).load()
+
+        next_id: int = len(workflow) + 1
         workflow_definition_id: str = _WORKFLOW_DEFINITION_ID_FORMAT.format(id=next_id)
-        self._workflow_definitions[workflow_definition_id] = workflow_definition
+        workflow[workflow_definition_id] = workflow_definition
+
+        with open(f"tests/{_WORKFLOW_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump(workflow)
+        UnitTestAPIAdapter.mp_lock.release()
+
         return {"id": workflow_definition_id}
 
     def get_workflow(self, *, workflow_definition_id: str) -> Dict[str, Any]:
-        if workflow_definition_id not in self._workflow_definitions:
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_WORKFLOW_PICKLE_FILE}", "rb") as pickle_file:
+            workflow = Unpickler(pickle_file).load()
+        UnitTestAPIAdapter.mp_lock.release()
+
+        if workflow_definition_id not in workflow:
             return {}
-        return {"workflow": self._workflow_definitions[workflow_definition_id]}
+        return {"workflow": workflow[workflow_definition_id]}
 
     def get_workflow_by_name(self, *, name: str, version: str) -> Dict[str, Any]:
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_WORKFLOW_PICKLE_FILE}", "rb") as pickle_file:
+            workflow = Unpickler(pickle_file).load()
+        UnitTestAPIAdapter.mp_lock.release()
+
         item = {}
-        for wfid, value in self._workflow_definitions.items():
+        for wfid, value in workflow.items():
             if value["name"] == name:
                 item = {"id": wfid, "workflow": value}
         return item
 
     def create_running_workflow(self, *, workflow_definition_id: str) -> str:
-        next_id: int = len(self._running_workflow) + 1
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_RUNNING_WORKFLOW_PICKLE_FILE}", "rb") as pickle_file:
+            running_workflow = Unpickler(pickle_file).load()
+
+        next_id: int = len(running_workflow) + 1
         running_workflow_id: str = _RUNNING_WORKFLOW_ID_FORMAT.format(id=next_id)
         record = {"done": False, "success": False, "workflow": workflow_definition_id}
-        self._running_workflow[running_workflow_id] = record
+        running_workflow[running_workflow_id] = record
+
+        with open(f"tests/{_RUNNING_WORKFLOW_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump(running_workflow)
+        UnitTestAPIAdapter.mp_lock.release()
+
         return {"id": running_workflow_id}
 
     def get_running_workflow(self, *, running_workflow_id: str) -> Dict[str, Any]:
-        if running_workflow_id not in self._running_workflow:
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_RUNNING_WORKFLOW_PICKLE_FILE}", "rb") as pickle_file:
+            running_workflow = Unpickler(pickle_file).load()
+        UnitTestAPIAdapter.mp_lock.release()
+
+        if running_workflow_id not in running_workflow:
             return {}
-        return {"running_workflow": self._running_workflow[running_workflow_id]}
+        return {"running_workflow": running_workflow[running_workflow_id]}
 
     def create_running_workflow_step(
         self, *, running_workflow_id: str, step: str
     ) -> str:
-        next_id: int = len(self._running_workflow_steps) + 1
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_RUNNING_WORKFLOW_STEP_PICKLE_FILE}", "rb") as pickle_file:
+            running_workflow_step = Unpickler(pickle_file).load()
+
+        next_id: int = len(running_workflow_step) + 1
         running_workflow_step_id: str = _RUNNING_WORKFLOW_STEP_ID_FORMAT.format(
             id=next_id
         )
@@ -78,58 +133,99 @@ def create_running_workflow_step(
             "success": False,
             "running_workflow": running_workflow_id,
         }
-        self._running_workflow_steps[running_workflow_step_id] = record
+        running_workflow_step[running_workflow_step_id] = record
+
+        with open(f"tests/{_RUNNING_WORKFLOW_STEP_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump(running_workflow_step)
+        UnitTestAPIAdapter.mp_lock.release()
+
         return {"id": running_workflow_step_id}
 
     def get_running_workflow_step(
         self, *, running_workflow_step_id: str
     ) -> Dict[str, Any]:
-        if running_workflow_step_id not in self._running_workflow_steps:
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_RUNNING_WORKFLOW_STEP_PICKLE_FILE}", "rb") as pickle_file:
+            running_workflow_step = Unpickler(pickle_file).load()
+        UnitTestAPIAdapter.mp_lock.release()
+
+        if running_workflow_step_id not in running_workflow_step:
             return {}
         return {
-            "running_workflow_step": self._running_workflow_steps[
-                running_workflow_step_id
-            ]
+            "running_workflow_step": running_workflow_step[running_workflow_step_id]
         }
 
     def get_running_workflow_steps(
         self, *, running_workflow_id: str
     ) -> List[Dict[str, Any]]:
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_RUNNING_WORKFLOW_STEP_PICKLE_FILE}", "rb") as pickle_file:
+            running_workflow_step = Unpickler(pickle_file).load()
+        UnitTestAPIAdapter.mp_lock.release()
+
         steps = []
-        for key, value in self._running_workflow_steps.items():
+        for key, value in running_workflow_step.items():
             if value["running_workflow"] == running_workflow_id:
                 item = {"running_workflow_step": value, "id": key}
                 steps.append(item)
         return {"count": len(steps), "running_workflow_steps": steps}
 
     def create_instance(self, *, running_workflow_step_id: str) -> Dict[str, Any]:
-        next_id: int = len(self._instances) + 1
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_INSTANCE_PICKLE_FILE}", "rb") as pickle_file:
+            instances = Unpickler(pickle_file).load()
+
+        next_id: int = len(instances) + 1
         instance_id: str = _INSTANCE_ID_FORMAT.format(id=next_id)
         record = {
             "running_workflow_step": running_workflow_step_id,
         }
-        self._instances[instance_id] = record
+        instances[instance_id] = record
+
+        with open(f"tests/{_INSTANCE_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump(instances)
+        UnitTestAPIAdapter.mp_lock.release()
+
         return {"id": instance_id}
 
     def get_instance(self, *, instance_id: str) -> Dict[str, Any]:
-        if instance_id not in self._instances:
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_INSTANCE_PICKLE_FILE}", "rb") as pickle_file:
+            instances = Unpickler(pickle_file).load()
+        UnitTestAPIAdapter.mp_lock.release()
+
+        if instance_id not in instances:
             return {}
-        return self._instances[instance_id]
+        return instances[instance_id]
 
     def create_task(self, *, instance_id: str) -> Dict[str, Any]:
-        next_id: int = len(self._instances) + 1
-        task_id: str = _INSTANCE_ID_FORMAT.format(id=next_id)
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_TASK_PICKLE_FILE}", "rb") as pickle_file:
+            tasks = Unpickler(pickle_file).load()
+
+        next_id: int = len(tasks) + 1
+        task_id: str = _TASK_ID_FORMAT.format(id=next_id)
         record = {
             "done": False,
             "exit_code": 0,
         }
-        self._tasks[task_id] = record
+        tasks[task_id] = record
+
+        with open(f"tests/{_TASK_PICKLE_FILE}", "wb") as pickle_file:
+            Pickler(pickle_file).dump(tasks)
+        UnitTestAPIAdapter.mp_lock.release()
+
         return {"id": task_id}
 
     def get_task(self, *, task_id: str) -> Dict[str, Any]:
-        if task_id not in self._tasks:
+        UnitTestAPIAdapter.mp_lock.acquire()
+        with open(f"tests/{_TASK_PICKLE_FILE}", "rb") as pickle_file:
+            tasks = Unpickler(pickle_file).load()
+        UnitTestAPIAdapter.mp_lock.release()
+
+        if task_id not in tasks:
             return {}
-        return self._tasks[task_id]
+        return tasks[task_id]
 
     def get_job(
         self, *, collection: str, job: str, version: str
diff --git a/tests/message_dispatcher.py b/tests/message_dispatcher.py
@@ -12,4 +12,5 @@ def __init__(self, msg_queue: UnitTestMessageQueue):
         self._msg_queue: UnitTestMessageQueue = msg_queue
 
     def send(self, message: Message) -> None:
+        print(f"UnitTestMessageDispatcher.send:\n{message}")
         self._msg_queue.put(message)
diff --git a/tests/message_queue.py b/tests/message_queue.py
@@ -18,6 +18,11 @@ def __init__(self, receiver: Optional[Callable[[Message], None]] = None):
         self._queue = Queue()
         self._receiver = receiver
 
+    def set_receiver(self, receiver: Callable[[Message], None]):
+        """Set or replace the receiver function, used unit-testing the WorkflowEngine."""
+        assert receiver
+        self._receiver = receiver
+
     def run(self):
         while not self._stop.is_set():
             with suppress(Empty):
diff --git a/tests/test_test_instance_launcher.py b/tests/test_test_instance_launcher.py
@@ -38,7 +38,7 @@ def test_get_nop_job(basic_launcher):
         workflow_id="workflow-00000000-0000-0000-0000-000000000000",
         running_workflow_step_id=rwfsid,
         workflow_definition={},
-        step_specification={"job": "nop"},
+        step_specification={"job": "nop", "variables": {"x": 1}},
     )
 
     # Assert
diff --git a/tests/test_workflow_engine_examples.py b/tests/test_workflow_engine_examples.py
@@ -1,8 +1,10 @@
 # A unit (functional) test for the WorkflowEngine's handling of 'Example 1'.
+import os
 import time
 from datetime import datetime, timezone
 
 import pytest
+import yaml
 
 pytestmark = pytest.mark.unit
 
@@ -23,33 +25,54 @@ def basic_engine():
     instance_launcher = UnitTestInstanceLauncher(
         api_adapter=api_adapter, msg_dispatcher=message_dispatcher
     )
+    workflow_engine = WorkflowEngine(
+        api_adapter=api_adapter, instance_launcher=instance_launcher
+    )
+    message_queue.set_receiver(workflow_engine.handle_message)
     return [
         api_adapter,
+        message_queue,
         message_dispatcher,
-        WorkflowEngine(api_adapter=api_adapter, instance_launcher=instance_launcher),
+        workflow_engine,
     ]
 
 
-def test_workflow_engine_with_example_1(basic_engine):
+def test_workflow_engine_with_two_step_nop(basic_engine):
     # Arrange
-    da, md, _ = basic_engine
-    # LOAD THE EXAMPLE-1 WORKFLOW DEFINITION INTO THE DATABASE
-    # TODO
-    # SIMULATE THE API CREATION OF A RUNNING WORKFLOW FROM THE WORKFLOW
-    wfid = da.create_workflow(workflow_definition={"name": "example-1"})
+    da, mq, md, _ = basic_engine
+
+    # Act
+    # To test the WorkflowEngine we need to:
+    # 1. Start the message queue
+    # 2. Load and create a Workflow Definition
+    # 3. Create a Running Workflow record
+    # 4. Send a Workflow START message
+    #
+    # 1. (Start the message queue)
+    mq.start()
+    # 2. (Load/create the workflow definition to be tested)
+    workflow_file_name = "example-two-step-nop"
+    workflow_path = os.path.join(
+        os.path.dirname(__file__), "workflow-definitions", f"{workflow_file_name}.yaml"
+    )
+    with open(workflow_path, "r", encoding="utf8") as wf_file:
+        wf_definition = yaml.load(wf_file, Loader=yaml.FullLoader)
+    assert wf_definition
+    wfid = da.create_workflow(workflow_definition=wf_definition)
     assert wfid
+    print(f"Created workflow definition {wfid}")
+    # 3. (Create a running workflow record)
     response = da.create_running_workflow(workflow_definition_id=wfid)
     r_wfid = response["id"]
     assert r_wfid
-
-    # Act
-    # SEND A MESSAGE TO THE ENGINE (VIA THE MESSAGE DISPATCHER) TO START THE WORKFLOW
-    # THE RUNNING WORKFLOW WILL HAVE THE ID "1"
+    print(f"Created running workflow {r_wfid}")
+    # 4. (Send the Workflow START message)
     msg = WorkflowMessage()
     msg.timestamp = f"{datetime.now(timezone.utc).isoformat()}Z"
     msg.action = "START"
     msg.running_workflow = r_wfid
     md.send(msg)
+    print("Sent START message")
 
     # Assert
     # Wait until the workflow is done (successfully)
@@ -59,6 +82,7 @@ def test_workflow_engine_with_example_1(basic_engine):
     r_wf = None
     while not done:
         response = da.get_running_workflow(running_workflow_id=r_wfid)
+        assert "running_workflow" in response
         r_wf = response["running_workflow"]
         if r_wf["done"]:
             done = True
@@ -67,14 +91,17 @@ def test_workflow_engine_with_example_1(basic_engine):
             if attempts > 10:
                 break
             time.sleep(0.5)
+    # Stop the message queue
+    print("Stopping message queue...")
+    mq.stop()
+    mq.join()
+    print("Stopped")
     assert r_wf
-    # TODO - The following should be 'success' but the implementation does not set it yet
-    assert not r_wf["success"]
+    assert r_wf["success"]
     # Now check there are the right number of RunningWorkflowStep Records
     # (and they're all set to success/done)
     response = da.get_running_workflow_steps(running_workflow_id=r_wfid)
-    # TODO - The following should not be zero but the implementation does not set it yet
-    assert response["count"] == 0
+    assert response["count"] == 2
     for step in response["running_workflow_steps"]:
         assert step["running_workflow_step"]["done"]
         assert step["running_workflow_step"]["success"]
diff --git a/tests/workflow-definitions/example-two-step-nop.yaml b/tests/workflow-definitions/example-two-step-nop.yaml
diff --git a/workflow/workflow_engine.py b/workflow/workflow_engine.py

Original file line number	Diff line number	Diff line change
`@@ -38,7 +38,7 @@ def test_get_nop_job(basic_launcher):`
`38`	`38`	`workflow_id="workflow-00000000-0000-0000-0000-000000000000",`
`39`	`39`	`running_workflow_step_id=rwfsid,`
`40`	`40`	`workflow_definition={},`
`41`		`- step_specification={"job": "nop"},`
	`41`	`+ step_specification={"job": "nop", "variables": {"x": 1}},`
`42`	`42`	`)`
`43`	`43`
`44`	`44`	`# Assert`