Synchronize repository with Athena deployment

2026-09-13 20:01:36 +02:00
parent 040a2df48b
commit fce9900389
60 changed files with 2570 additions and 607 deletions
@@ -26,7 +26,48 @@ def image_item(state="exited"):
            "Labels": {controller.IMAGE_LABEL_KEY: controller.IMAGE_WORKER}}


+def restore_item(state="exited"):
+    return {"Id": "id-restore", "State": state,
+            "Labels": {controller.IMAGE_LABEL_KEY: controller.RESTORE_WORKER}}
+
+
+def tts_item(state="running"):
+    return {"Id": "id-tts", "State": state,
+            "Labels": {controller.TTS_LABEL_KEY: controller.TTS_WORKER}}
+
+
+def music_item(state="exited"):
+    return {"Id": "id-music", "State": state,
+            "Labels": {controller.MUSIC_LABEL_KEY: "acestep"}}
+
+
 class ProfileControllerTests(unittest.TestCase):
+    def test_music_start_exclusively_stops_gpu_workers(self):
+        profiles = {name: item(name) for name in controller.ALLOWED}
+        profiles["ultra"] = item("ultra", "running")
+        calls = []
+
+        def request(method, path):
+            calls.append((method, path))
+            return 204, b""
+
+        with patch.object(controller, "MUSIC_WORKER", "acestep"), \
+             patch.object(controller, "containers", return_value=profiles), \
+             patch.object(controller, "music_container", return_value=music_item()), \
+             patch.object(controller, "image_containers",
+                          return_value=[image_item("running")]), \
+             patch.object(controller, "tts_container", return_value=tts_item()), \
+             patch.object(controller, "docker_request", side_effect=request):
+            result = controller.set_music_worker(True)
+
+        self.assertEqual(result, {"music_worker": "acestep", "state": "running"})
+        self.assertEqual(calls, [
+            ("POST", "/containers/id-ultra/stop?t=120"),
+            ("POST", "/containers/id-flux/stop?t=20"),
+            ("POST", "/containers/id-tts/stop?t=30"),
+            ("POST", "/containers/id-music/start"),
+        ])
+
    def test_rejects_unknown_profile_before_docker_call(self):
        with patch.object(controller, "docker_request") as request:
            with self.assertRaises(ValueError):
@@ -43,7 +84,8 @@ class ProfileControllerTests(unittest.TestCase):
            return 204, b""

        with patch.object(controller, "containers", return_value=profiles), \
-             patch.object(controller, "image_container", return_value=image_item()), \
+             patch.object(controller, "image_containers", return_value=[image_item()]), \
+             patch.object(controller, "tts_container", return_value=tts_item()), \
             patch.object(controller, "docker_request", side_effect=request):
            result = controller.activate("medium")

@@ -56,7 +98,8 @@ class ProfileControllerTests(unittest.TestCase):
    def test_fails_if_profile_container_is_missing(self):
        profiles = {name: item(name) for name in controller.ALLOWED[:-1]}
        with patch.object(controller, "containers", return_value=profiles), \
-             patch.object(controller, "image_container", return_value=image_item()):
+             patch.object(controller, "image_containers", return_value=[image_item()]), \
+             patch.object(controller, "tts_container", return_value=tts_item()):
            with self.assertRaisesRegex(RuntimeError, "missing"):
                controller.activate("fast")

@@ -71,13 +114,38 @@ class ProfileControllerTests(unittest.TestCase):

        with patch.object(controller, "containers", return_value=profiles), \
             patch.object(controller, "image_container", return_value=image_item()), \
+             patch.object(controller, "image_containers",
+                          return_value=[image_item(), restore_item()]), \
+             patch.object(controller, "tts_container", return_value=tts_item()), \
             patch.object(controller, "docker_request", side_effect=request):
            controller.set_image_worker(True)
        self.assertEqual(calls, [
            ("POST", "/containers/id-medium/stop?t=120"),
+            ("POST", "/containers/id-tts/stop?t=30"),
            ("POST", "/containers/id-flux/start"),
        ])

+    def test_restore_start_stops_flux_and_starts_restore(self):
+        profiles = {name: item(name) for name in controller.ALLOWED}
+        calls = []
+
+        def request(method, path):
+            calls.append((method, path))
+            return 204, b""
+
+        with patch.object(controller, "containers", return_value=profiles), \
+             patch.object(controller, "image_container", return_value=restore_item()), \
+             patch.object(controller, "image_containers",
+                          return_value=[image_item("running"), restore_item()]), \
+             patch.object(controller, "tts_container", return_value=tts_item()), \
+             patch.object(controller, "docker_request", side_effect=request):
+            controller.set_image_worker(True, controller.RESTORE_WORKER)
+        self.assertEqual(calls, [
+            ("POST", "/containers/id-tts/stop?t=30"),
+            ("POST", "/containers/id-flux/stop?t=20"),
+            ("POST", "/containers/id-restore/start"),
+        ])
+
    def test_profile_activation_stops_image_worker_first(self):
        profiles = {name: item(name) for name in controller.ALLOWED}
        calls = []
@@ -87,8 +155,9 @@ class ProfileControllerTests(unittest.TestCase):
            return 204, b""

        with patch.object(controller, "containers", return_value=profiles), \
-             patch.object(controller, "image_container",
-                          return_value=image_item("running")), \
+             patch.object(controller, "image_containers",
+                          return_value=[image_item("running"), restore_item()]), \
+             patch.object(controller, "tts_container", return_value=tts_item()), \
             patch.object(controller, "docker_request", side_effect=request):
            controller.activate("fast")
        self.assertEqual(calls, [