Get subprocess test to run

Author: AlanCoding

ansible_base/task/__init__.py

@@ -0,0 +1,19 @@
+from ansible_base.lib.utils.db import get_pg_notify_params
+
+
+def get_config():
+    psycopg_params = get_pg_notify_params()
+    psycopg_params.pop('autocommit')  # dispatcher automatically adds this, causes error, TODO: need pre-check
+    psycopg_params.pop('cursor_factory')
+    psycopg_params.pop('context')  # TODO: remove in inner method, makes non-async, not good
+
+    return {
+        "producers": {
+            "brokers": {
+                "pg_notify": psycopg_params,
+                "channels": ["dab_broadcast"],
+            },
+            "scheduled": {},
+        },
+        "pool": {"max_workers": 4},  # TODO: to settings
+    }

ansible_base/task/management/commands/run_dispatcher.py

@@ -3,7 +3,7 @@
 
 from dispatcher.main import DispatcherMain
 
-from ansible_base.lib.utils.db import get_pg_notify_params
+from ansible_base.task import get_config
 
 from django.core.management.base import BaseCommand
 from django.db import connection
@@ -16,21 +16,7 @@ class Command(BaseCommand):
     help = "Runs bug checking sanity checks, gets scale metrics, and recommendations for Role Based Access Control"
 
     def handle(self, *args, **options):
-        psycopg_params = get_pg_notify_params()
-        psycopg_params.pop('autocommit')  # dispatcher automatically adds this, causes error, TODO: need pre-check
-        psycopg_params.pop('cursor_factory')
-        psycopg_params.pop('context')  # TODO: remove in inner method, makes non-async, not good
-
-        dispatcher_config = {
-            "producers": {
-                "brokers": {
-                    "pg_notify": psycopg_params,
-                    "channels": ["dab_broadcast"],
-                },
-                "scheduled": {},
-            },
-            "pool": {"max_workers": 4},  # TODO: to settings
-        }
+        dispatcher_config = get_config()
 
         loop = asyncio.get_event_loop()
         dispatcher = DispatcherMain(dispatcher_config)

test_app/tasks.py

@@ -1,4 +1,7 @@
 import time
+from uuid import UUID
+
+from test_app.models import UUIDModel
 
 from ansible_base.task.publish import durable_task
 
@@ -8,6 +11,11 @@ def hello_world():
     print('hello world')
 
 
+@durable_task()
+def create_uuid_entry(uuid: UUID):
+    UUIDModel.objects.create(id=uuid)
+
+
 @durable_task()
 def sleep(seconds=2):
     print(f'about to sleep for {seconds} seconds')

test_app/tests/task/conftest.py

@@ -0,0 +1,87 @@
+import asyncio
+import multiprocessing
+import contextlib
+import time
+
+import pytest
+
+
+from dispatcher.main import DispatcherMain
+
+
+
+async def asyncio_target(queue_in, queue_out, config):
+    try:
+        dispatcher = DispatcherMain(config)
+
+        await dispatcher.connect_signals()
+        await dispatcher.start_working()
+        await dispatcher.wait_for_producers_ready()
+        queue_out.put('ready')
+
+
+        print('dispatcher server listening on queue_in')
+        loop = asyncio.get_event_loop()
+        message = await loop.run_in_executor(None, queue_in.get)
+
+        print(f'got message, will shut down: {message}')
+    finally:
+        await dispatcher.shutdown()
+        await dispatcher.cancel_tasks()
+
+
+def subprocess_target(queue_in, queue_out, config):
+    loop = asyncio.get_event_loop()
+    try:
+        loop.run_until_complete(asyncio_target(queue_in, queue_out, config))
+    except Exception:
+        import traceback
+
+        traceback.print_exc()
+        # We are in a subprocess here, so even if we handle the exception
+        # the main process will not know and still wait forever
+        # so give them a kick on our way out
+        print('sending error message after error')
+        queue_out.put('error')
+    finally:
+        print('closing asyncio loop')
+        loop.close()
+
+
+class SubprocessRunner:
+
+    def __init__(self):
+        self.queue_in = multiprocessing.Queue()
+        self.queue_out = multiprocessing.Queue()
+
+    def start_in_subprocess(self, config):
+        process = multiprocessing.Process(target=subprocess_target, args=(self.queue_in, self.queue_out, config))
+        process.start()
+        return process
+
+    @contextlib.contextmanager
+    def with_server(self, config):
+        process = self.start_in_subprocess(config)
+        msg = self.queue_out.get()
+        if msg != 'ready':
+            raise RuntimeError('never got ready message from subprocess')
+        try:
+            yield self
+        finally:
+            self.queue_in.put('stop')
+            process.terminate()  # SIGTERM
+            # Poll to close process resources, due to race condition where it is not still running
+            for i in range(3):
+                time.sleep(0.1)
+                try:
+                    process.close()
+                    break
+                except Exception:
+                    if i == 2:
+                        raise
+
+
+@pytest.fixture
+def dispatcher_subprocess():
+    server = SubprocessRunner()
+    return server.with_server

test_app/tests/task/test_basic.py

@@ -0,0 +1,23 @@
+from uuid import uuid4
+import time
+
+import pytest
+
+from ansible_base.task import get_config
+
+from test_app.tasks import create_uuid_entry
+from test_app.models import UUIDModel
+
+
+@pytest.mark.django_db
+def test_run_task(dispatcher_subprocess):
+    with dispatcher_subprocess(get_config()) as server:
+        my_uuid = str(uuid4())
+        create_uuid_entry.delay(uuid=my_uuid)
+
+        for i in range(20):
+            if UUIDModel.objects.filter(id=my_uuid).exists():
+                break
+            time.sleep(0.05)
+        else:
+            assert f'Task never inserted UUID entry with {my_uuid}'