ss

d2ypp2/tests/test_neko_capacity.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from src.neko_capacity import (
+    desired_neko_count_for_youtube_channels,
+    ensure_neko_capacity_for_youtube_channels,
+    upsert_env_value,
+)
+
+
+class NekoCapacityTests(unittest.TestCase):
+    def test_desired_count_keeps_two_instance_floor_and_sixty_four_instance_cap(self) -> None:
+        self.assertEqual(desired_neko_count_for_youtube_channels([]), 2)
+        self.assertEqual(desired_neko_count_for_youtube_channels([{"id": "lane1"}]), 2)
+        self.assertEqual(desired_neko_count_for_youtube_channels([{"id": "lane1"}, {"id": "lane2"}]), 2)
+        self.assertEqual(
+            desired_neko_count_for_youtube_channels(
+                [{"id": "lane1"}, {"id": "lane2"}, {"id": "lane3"}, {"id": "lane4"}]
+            ),
+            4,
+        )
+        self.assertEqual(desired_neko_count_for_youtube_channels({"id": f"lane{i}"} for i in range(80)), 64)
+
+    def test_upsert_env_value_preserves_file_and_changes_target_key(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            env = Path(tmp) / "config" / "system-stack.env"
+            env.parent.mkdir()
+            env.write_text("HOSTNAME_ALIAS=live\nNEKO_INSTANCE_COUNT=2\nPORT=8001\n", encoding="utf-8")
+
+            changed = upsert_env_value(env, "NEKO_INSTANCE_COUNT", "3")
+
+            self.assertTrue(changed)
+            self.assertEqual(env.read_text(encoding="utf-8"), "HOSTNAME_ALIAS=live\nNEKO_INSTANCE_COUNT=3\nPORT=8001\n")
+
+    def test_ensure_capacity_increases_env_and_restarts_neko_in_background(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            env = root / "config" / "system-stack.env"
+            env.parent.mkdir()
+            env.write_text("NEKO_INSTANCE_COUNT=2\n", encoding="utf-8")
+            service_ctl = root / "scripts" / "linux" / "service_ctl.sh"
+            service_ctl.parent.mkdir(parents=True)
+            service_ctl.write_text("#!/usr/bin/env bash\n", encoding="utf-8")
+
+            with patch("src.neko_capacity.subprocess.Popen") as popen:
+                result = ensure_neko_capacity_for_youtube_channels(
+                    root,
+                    [{"id": "lane1"}, {"id": "lane2"}, {"id": "lane3"}],
+                )
+
+            self.assertTrue(result["changed"])
+            self.assertTrue(result["restarted"])
+            self.assertIn("NEKO_INSTANCE_COUNT=3", env.read_text(encoding="utf-8"))
+            popen.assert_called_once()
+            self.assertEqual(popen.call_args.args[0][-2:], ["neko", "restart"])
+
+    def test_ensure_capacity_keeps_saved_env_when_restart_spawn_fails(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            env = root / "config" / "system-stack.env"
+            env.parent.mkdir()
+            env.write_text("NEKO_INSTANCE_COUNT=2\n", encoding="utf-8")
+            service_ctl = root / "scripts" / "linux" / "service_ctl.sh"
+            service_ctl.parent.mkdir(parents=True)
+            service_ctl.write_text("#!/usr/bin/env bash\n", encoding="utf-8")
+
+            with patch("src.neko_capacity.subprocess.Popen", side_effect=OSError("fork failed")):
+                result = ensure_neko_capacity_for_youtube_channels(
+                    root,
+                    [{"id": "lane1"}, {"id": "lane2"}, {"id": "lane3"}],
+                )
+
+            self.assertTrue(result["changed"])
+            self.assertFalse(result["restarted"])
+            self.assertIn("restart_spawn_failed", result["reason"])
+            self.assertIn("NEKO_INSTANCE_COUNT=3", env.read_text(encoding="utf-8"))
+
+    def test_ensure_capacity_does_not_downscale_manual_larger_count(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            env = root / "config" / "system-stack.env"
+            env.parent.mkdir()
+            env.write_text("NEKO_INSTANCE_COUNT=3\n", encoding="utf-8")
+
+            result = ensure_neko_capacity_for_youtube_channels(root, [{"id": "lane1"}], restart=False)
+
+            self.assertFalse(result["changed"])
+            self.assertEqual(env.read_text(encoding="utf-8"), "NEKO_INSTANCE_COUNT=3\n")
+
+
+if __name__ == "__main__":
+    unittest.main()