Manage graph viewers with a supervised local service

tests/test_visualization.py

@@ -3,22 +3,23 @@ from __future__ import annotations
 import json
 import shutil
 import subprocess
-import sys
 import tempfile
+import threading
 import time
 import unittest
 import urllib.error
 import urllib.parse
 import urllib.request
+from contextlib import contextmanager
 from pathlib import Path
 
 from docforge.errors import DocForgeError
 from docforge.index import ProjectIndex
 from docforge.project import Project
+from docforge.viewer_manager import ViewerManager, ViewerManagerClient
 from docforge.visualization import (
     _GRAPH_BROWSER_HTML,
     VISUALIZATION_TEMPLATE,
-    PersistentVisualizationRunner,
     VisualizationIndexSnapshot,
     VisualizationRunner,
 )
@@ -33,6 +34,31 @@ class VisualizationTests(unittest.TestCase):
         shutil.copytree(FIXTURES / name, root)
         return root
 
+    @contextmanager
+    def running_manager(
+        self,
+        state_path: Path,
+        *,
+        idle_timeout_seconds: float = 60,
+        check_interval_seconds: float = 0.02,
+    ):
+        manager = ViewerManager(
+            state_path,
+            idle_timeout_seconds=idle_timeout_seconds,
+            check_interval_seconds=check_interval_seconds,
+        )
+        thread = threading.Thread(target=manager.serve_forever, daemon=True)
+        thread.start()
+        deadline = time.monotonic() + 2
+        while not state_path.exists() and time.monotonic() < deadline:
+            time.sleep(0.01)
+        self.assertTrue(state_path.exists())
+        try:
+            yield manager
+        finally:
+            manager.shutdown()
+            thread.join(timeout=2)
+
     def test_overview_and_neighborhood_are_deterministic_and_bounded(self) -> None:
         with tempfile.TemporaryDirectory() as directory:
             root = self.copy_fixture("alpha", Path(directory))
@@ -380,52 +406,46 @@ if (dependencyEdge.source_id !== "dependency" || dependencyEdge.target_id !== "p
             finally:
                 runner.stop()
 
-    def test_persistent_worker_survives_launcher_and_stops_only_explicitly(self) -> None:
+    def test_manager_reuses_workers_and_applies_explicit_or_idle_shutdown(self) -> None:
         with tempfile.TemporaryDirectory() as directory:
             root = self.copy_fixture("alpha", Path(directory))
-            ProjectIndex(Project.open(root)).build()
-            script = """
-import sys
-from pathlib import Path
-from docforge.index import ProjectIndex
-from docforge.project import Project
-from docforge.visualization import PersistentVisualizationRunner
+            index = ProjectIndex(Project.open(root))
+            index.build()
+            state_path = Path(directory) / "viewer-manager.json"
+            with self.running_manager(
+                state_path, idle_timeout_seconds=0.25, check_interval_seconds=0.02
+            ):
+                client = ViewerManagerClient(index, state_path=state_path)
+                first = client.start()
+                url = str(first["url"])
+                self.assertEqual("managed_idle", first["lifetime"]["policy"])
+                self.assertTrue(url.startswith("http://127.0.0.1:"))
+                with urllib.request.urlopen(url, timeout=2) as response:
+                    self.assertEqual(200, response.status)
 
-runner = PersistentVisualizationRunner(ProjectIndex(Project.open(Path(sys.argv[1]))))
-print(runner.start()["url"], flush=True)
-time.sleep(60)
-"""
-            script = "import time\n" + script
-            with subprocess.Popen(
-                [sys.executable, "-c", script, str(root)],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                text=True,
-            ) as launcher:
-                assert launcher.stdout is not None
-                url = launcher.stdout.readline().strip()
-                launcher.terminate()
-                launcher.wait(timeout=2)
-                self.assertLess(launcher.returncode, 0)
-            self.assertTrue(url.startswith("http://127.0.0.1:"))
-            with urllib.request.urlopen(url, timeout=2) as response:
-                self.assertEqual(200, response.status)
-
-            time.sleep(0.6)
-            with urllib.request.urlopen(url, timeout=2) as response:
-                self.assertEqual(200, response.status)
-
-            runner = PersistentVisualizationRunner(ProjectIndex(Project.open(root)))
-            try:
-                reused = runner.start()
+                reused = client.start(depth=2)
                 self.assertTrue(reused["reused"])
                 self.assertEqual(url.split("?", 1)[0], str(reused["url"]).split("?", 1)[0])
-                stopped = runner.stop()
+                self.assertEqual("running", client.status()["state"])
+
+                time.sleep(0.1)
+                heartbeat = url.split("?", 1)[0] + "api/heartbeat"
+                with urllib.request.urlopen(heartbeat, timeout=2) as response:
+                    self.assertEqual("alive", json.load(response)["viewer"])
+                time.sleep(0.1)
+                self.assertEqual("running", client.status()["state"])
+
+                time.sleep(0.35)
+                deadline = time.monotonic() + 2
+                while client.status()["state"] == "running" and time.monotonic() < deadline:
+                    time.sleep(0.02)
+                self.assertEqual("not_running", client.status()["state"])
+
+                restarted = client.start()
+                stopped = client.stop()
                 self.assertEqual("stopped", stopped["state"])
                 with self.assertRaises(OSError):
-                    urllib.request.urlopen(url, timeout=0.2)
-            finally:
-                runner.stop()
+                    urllib.request.urlopen(str(restarted["url"]), timeout=0.2)
 
     def test_runner_rejects_ambiguous_targets_and_changed_index_snapshot(self) -> None:
         with tempfile.TemporaryDirectory() as directory: