Fix test run w/ Go server/worker

Author: cretz

.gitignore

@@ -6,3 +6,5 @@ temporalio/api/*
 temporalio/bridge/proto/*
 !temporalio/bridge/proto/__init__.py
 temporalio/bridge/target/
+/tests/fixtures/golangserver/golangserver
+/tests/fixtures/golangworker/golangworker

poetry.lock

@@ -24,7 +24,6 @@ isort = "^5.10.1"
 maturin = "^0.12.6"
 mypy = "^0.931"
 mypy-protobuf = "^3.2.0"
-psutil = "^5.9.0"
 pydocstyle = "^6.1.1"
 pytest = "^6.2.5"
 pytest-asyncio = "^0.18.0"

pyproject.toml

@@ -9,8 +9,6 @@
 def event_loop():
     # See https://github.com/pytest-dev/pytest-asyncio/issues/68
     # See https://github.com/pytest-dev/pytest-asyncio/issues/257
-    policy = asyncio.get_event_loop_policy()
-    loop = policy.new_event_loop()
+    loop = asyncio.get_event_loop_policy().new_event_loop()
     yield loop
     loop.close()
-    policy.set_event_loop(loop)

tests/conftest.py

@@ -1,13 +1,12 @@
 import asyncio
 import logging
 import os
-import signal
+import subprocess
 import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Any, AsyncGenerator, Optional
 
-import psutil
 import pytest_asyncio
 
 import temporalio.api.workflowservice.v1
@@ -60,15 +59,11 @@ async def start() -> "ExternalGolangServer":
         namespace = f"test-namespace-{uuid.uuid4()}"
         # TODO(cretz): Make this configurable?
         port = "9233"
-        process = await asyncio.create_subprocess_exec(
-            "go",
-            "run",
-            ".",
+        process = await start_external_go_process(
+            os.path.join(os.path.dirname(__file__), "golangserver"),
+            "golangserver",
             port,
             namespace,
-            cwd=os.path.abspath(
-                os.path.join(os.path.dirname(__file__), "golangserver")
-            ),
         )
         server = ExternalGolangServer(f"localhost:{port}", namespace, process)
         # Try to get the client multiple times to check whether server is online
@@ -99,7 +94,7 @@ def namespace(self) -> str:
         return self._namespace
 
     async def close(self):
-        kill_proc_tree(self._process.pid)
+        self._process.terminate()
         await self._process.wait()
 
 
@@ -143,19 +138,15 @@ async def close(self):
 class ExternalGolangWorker(Worker):
     @staticmethod
     async def start(host_port: str, namespace: str) -> "ExternalGolangWorker":
-        task_queue = uuid.uuid4()
-        process = await asyncio.create_subprocess_exec(
-            "go",
-            "run",
-            ".",
+        task_queue = str(uuid.uuid4())
+        process = await start_external_go_process(
+            os.path.join(os.path.dirname(__file__), "golangworker"),
+            "golangworker",
             host_port,
             namespace,
-            str(task_queue),
-            cwd=os.path.abspath(
-                os.path.join(os.path.dirname(__file__), "golangworker")
-            ),
+            task_queue,
         )
-        return ExternalGolangWorker(str(task_queue), process)
+        return ExternalGolangWorker(task_queue, process)
 
     def __init__(self, task_queue: str, process: asyncio.subprocess.Process) -> None:
         super().__init__()
@@ -167,7 +158,7 @@ def task_queue(self) -> str:
         return self._task_queue
 
     async def close(self):
-        kill_proc_tree(self._process.pid)
+        self._process.terminate()
         await self._process.wait()
 
 
@@ -178,24 +169,14 @@ async def worker(server: Server) -> AsyncGenerator[Worker, None]:
     await worker.close()
 
 
-# See https://psutil.readthedocs.io/en/latest/#kill-process-tree
-def kill_proc_tree(
-    pid, sig=signal.SIGTERM, include_parent=True, timeout=None, on_terminate=None
-):
-    """Kill a process tree (including grandchildren) with signal
-    "sig" and return a (gone, still_alive) tuple.
-    "on_terminate", if specified, is a callback function which is
-    called as soon as a child terminates.
-    """
-    assert pid != os.getpid(), "won't kill myself"
-    parent = psutil.Process(pid)
-    children = parent.children(recursive=True)
-    if include_parent:
-        children.append(parent)
-    for p in children:
-        try:
-            p.send_signal(sig)
-        except psutil.NoSuchProcess:
-            pass
-    gone, alive = psutil.wait_procs(children, timeout=timeout, callback=on_terminate)
-    return (gone, alive)
+async def start_external_go_process(
+    source_dir: str, exe_name: str, *args: str
+) -> asyncio.subprocess.Process:
+    # First, build the executable. We accept the performance issues of building
+    # this each run.
+    logger.info("Building %s", exe_name)
+    subprocess.run(["go", "build", "-o", exe_name, "."], cwd=source_dir, check=True)
+    logger.info("Starting %s", exe_name)
+    return await asyncio.create_subprocess_exec(
+        os.path.join(source_dir, exe_name), *args
+    )