190 lines
8.8 KiB
Python
190 lines
8.8 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from subprocess import CompletedProcess
|
|
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_synchronously(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._compose_has_neko_service", side_effect=[False, True]):
|
|
with patch(
|
|
"src.neko_capacity.subprocess.run",
|
|
return_value=CompletedProcess(args=["service_ctl"], returncode=0),
|
|
) as run:
|
|
result = ensure_neko_capacity_for_youtube_channels(
|
|
root,
|
|
[{"id": "lane1"}, {"id": "lane2"}, {"id": "lane3"}],
|
|
)
|
|
|
|
self.assertTrue(result["ok"])
|
|
self.assertTrue(result["changed"])
|
|
self.assertTrue(result["restart_required"])
|
|
self.assertTrue(result["restarted"])
|
|
self.assertTrue(result["compose_ready"])
|
|
self.assertIn("NEKO_INSTANCE_COUNT=3", env.read_text(encoding="utf-8"))
|
|
run.assert_called_once()
|
|
self.assertEqual(run.call_args.args[0][-2:], ["neko", "restart"])
|
|
|
|
def test_ensure_capacity_repairs_stale_generated_compose_even_when_env_count_is_high_enough(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")
|
|
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._compose_has_neko_service", side_effect=[False, True]):
|
|
with patch(
|
|
"src.neko_capacity.subprocess.run",
|
|
return_value=CompletedProcess(args=["service_ctl"], returncode=0),
|
|
) as run:
|
|
result = ensure_neko_capacity_for_youtube_channels(
|
|
root,
|
|
[{"id": "lane1"}, {"id": "lane2"}, {"id": "lane3"}],
|
|
)
|
|
|
|
self.assertTrue(result["ok"])
|
|
self.assertFalse(result["changed"])
|
|
self.assertTrue(result["restart_required"])
|
|
self.assertTrue(result["restarted"])
|
|
self.assertIn("compose_missing_desired_service", result["reason"])
|
|
run.assert_called_once()
|
|
|
|
def test_ensure_capacity_reports_restart_success_but_missing_compose_as_not_ok(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._compose_has_neko_service", return_value=False):
|
|
with patch(
|
|
"src.neko_capacity.subprocess.run",
|
|
return_value=CompletedProcess(args=["service_ctl"], returncode=0),
|
|
):
|
|
result = ensure_neko_capacity_for_youtube_channels(
|
|
root,
|
|
[{"id": "lane1"}, {"id": "lane2"}, {"id": "lane3"}],
|
|
)
|
|
|
|
self.assertFalse(result["ok"])
|
|
self.assertTrue(result["restarted"])
|
|
self.assertFalse(result["compose_ready"])
|
|
self.assertIn("compose_missing_desired_service_after_restart", result["reason"])
|
|
|
|
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._compose_has_neko_service", return_value=False):
|
|
with patch("src.neko_capacity.subprocess.run", 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.assertFalse(result["ok"])
|
|
self.assertIn("restart_failed", result["reason"])
|
|
self.assertIn("NEKO_INSTANCE_COUNT=3", env.read_text(encoding="utf-8"))
|
|
|
|
def test_ensure_capacity_keeps_saved_env_when_restart_exits_nonzero(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._compose_has_neko_service", return_value=False):
|
|
with patch(
|
|
"src.neko_capacity.subprocess.run",
|
|
return_value=CompletedProcess(args=["service_ctl"], returncode=17),
|
|
):
|
|
result = ensure_neko_capacity_for_youtube_channels(
|
|
root,
|
|
[{"id": "lane1"}, {"id": "lane2"}, {"id": "lane3"}],
|
|
)
|
|
|
|
self.assertTrue(result["changed"])
|
|
self.assertFalse(result["restarted"])
|
|
self.assertFalse(result["ok"])
|
|
self.assertIn("restart_exit_17", 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")
|
|
|
|
with patch("src.neko_capacity._compose_has_neko_service", return_value=True):
|
|
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()
|