import asyncio import contextlib import io import json from pathlib import Path import sys import tempfile from types import SimpleNamespace import unittest from unittest.mock import AsyncMock, Mock, patch ROOT = Path(__file__).resolve().parents[1] sys.path[:0] = [str(ROOT), str(ROOT / "vendor")] from awg_gus import manage from awg_gus.bot import Bot from awg_gus.core import State from awg_gus.main import configure, load_config, save_config from awg_gus.profiles import Layout, assert_unique_bot, profile_name from test_bot import FakeAwg, FakeDelivery, OWNER, OTHER, callback, message BOT_A = "aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa" BOT_B = "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb" def config(layout, name, bot=BOT_A): return dict(url=f"https://{name}.example.test", token="gvb_test_" + bot, bot_id=bot, owners=[OWNER], awg2=str(layout.root / "bin/awg2"), state_dir=str(layout.state(name))) class FakeSystemd: def __init__(self): self.calls = [] self.active = {Layout.unit("default")} self.enabled = set(self.active) self.fail_restart = False def __call__(self, *args, check=True): self.calls.append(args) action, unit = args[0], args[-1] rc = 0 if action == "is-active": rc = 0 if unit in self.active else 3 elif action == "is-enabled": rc = 0 if unit in self.enabled else 1 elif action == "stop": self.active.discard(unit) elif action == "disable": self.enabled.discard(unit) elif action == "enable": self.enabled.add(unit) elif action in {"start", "restart"}: if action == "restart" and self.fail_restart: raise RuntimeError("simulated startup failure") self.active.add(unit) return SimpleNamespace(returncode=rc, stdout="active" if rc == 0 else "inactive") class ProfileTests(unittest.TestCase): def setUp(self): self.tmp = tempfile.TemporaryDirectory() self.layout = Layout(Path(self.tmp.name)) save_config(self.layout.config("default"), config(self.layout, "default")) def tearDown(self): self.tmp.cleanup() def test_legacy_paths_preserved_and_new_paths_isolated(self): self.assertEqual(self.layout.names(), ["default"]) self.assertEqual(self.layout.config("default"), self.layout.root / "etc/awg-gus-bot.json") self.assertNotEqual(self.layout.state("gus2"), self.layout.state("default")) for name in ["../bad", "x/y", "bad;id", "a.service", "-name", "A", "a"*33]: with self.subTest(name=name), self.assertRaises(ValueError): profile_name(name) def test_duplicate_bot_rejected_even_with_alias_url_and_rotated_token(self): candidate = config(self.layout, "gus2") candidate["token"] = "gvb_rotated" with self.assertRaisesRegex(ValueError, "default"): assert_unique_bot(self.layout.config("gus2"), candidate, self.layout) def test_second_bot_and_existing_token_rotation_are_allowed(self): candidate = config(self.layout, "gus2", BOT_B) assert_unique_bot(self.layout.config("gus2"), candidate, self.layout) original = config(self.layout, "default") original["token"] = "gvb_rotated" assert_unique_bot(self.layout.config("default"), original, self.layout) def test_existing_profile_cannot_silently_switch_identity(self): candidate = config(self.layout, "default", BOT_B) with self.assertRaisesRegex(ValueError, "другой бот"): assert_unique_bot(self.layout.config("default"), candidate, self.layout) def test_duplicate_is_rejected_before_pairing_or_polling(self): client = Mock() client.get_me.return_value = dict(id=BOT_A, api_version=3, features=["commands", "keyboards", "parse_mode"]) client.get_webhook_info.return_value = {} client.api_url.return_value = "https://alias.example.test" args = SimpleNamespace(config=str(self.layout.config("gus2")), layout=self.layout) with patch("awg_gus.main.sys.stdin.isatty", return_value=True), \ patch("builtins.input", return_value="https://alias.example.test"), \ patch("awg_gus.main.getpass.getpass", return_value="gvb_rotated"), \ patch("awg_gus.main.Client", return_value=client), \ patch("awg_gus.main.pair") as pair, contextlib.redirect_stdout(io.StringIO()): with self.assertRaisesRegex(ValueError, "default"): configure(args) pair.assert_not_called() client.get_updates.assert_not_called() self.assertFalse(self.layout.config("gus2").exists()) def test_add_starts_only_new_unit_and_preserves_legacy_config(self): original = self.layout.config("default").read_bytes() systemd = FakeSystemd() def setup(args): save_config(args.config, config(self.layout, "gus2", BOT_B)) with patch("awg_gus.manage.systemctl", systemd), \ patch("awg_gus.manage.shutil.which", return_value=str(self.layout.root / "bin/awg2")), \ patch("awg_gus.manage.Awg.call", new_callable=AsyncMock), \ patch("awg_gus.manage.time.sleep"), contextlib.redirect_stdout(io.StringIO()): manage.setup_profile("gus2", layout=self.layout, setup_fn=setup, check_fn=AsyncMock()) self.assertEqual(self.layout.config("default").read_bytes(), original) self.assertEqual(self.layout.names(), ["default", "gus2"]) self.assertEqual(systemd.active, {Layout.unit("default"), Layout.unit("gus2")}) self.assertFalse(any(c[-1] == Layout.unit("default") for c in systemd.calls)) template = (self.layout.root / "etc/systemd/system/awg-gus-bot@.service").read_text() self.assertIn("/instances/%i.json", template) self.assertIn("Restart=on-failure", template) def test_failed_add_rolls_back_only_new_profile(self): original = self.layout.config("default").read_bytes() systemd = FakeSystemd() systemd.fail_restart = True def setup(args): save_config(args.config, config(self.layout, "gus2", BOT_B)) s = State(Path(args.state_dir) / "state.sqlite3", "new") s.set("offset", 22) s.db.close() with patch("awg_gus.manage.systemctl", systemd), \ patch("awg_gus.manage.shutil.which", return_value=str(self.layout.root / "bin/awg2")), \ patch("awg_gus.manage.Awg.call", new_callable=AsyncMock): with self.assertRaises(RuntimeError): manage.setup_profile("gus2", layout=self.layout, setup_fn=setup, check_fn=AsyncMock()) self.assertEqual(self.layout.config("default").read_bytes(), original) self.assertFalse(self.layout.config("gus2").exists()) self.assertFalse((self.layout.state("gus2") / "state.sqlite3").exists()) self.assertEqual(systemd.active, {Layout.unit("default")}) self.assertNotIn(Layout.unit("gus2"), systemd.enabled) def test_failed_reconfigure_restores_cursor_config_and_service(self): original = self.layout.config("default").read_bytes() database = self.layout.state("default") / "state.sqlite3" s = State(database, "original") s.set("offset", 91) s.db.close() systemd = FakeSystemd() def setup(args): changed = config(self.layout, "default") changed["token"] = "gvb_changed" save_config(args.config, changed) s = State(database, "original") s.set("offset", 999) s.db.close() raise ValueError("pairing failed") with patch("awg_gus.manage.systemctl", systemd), \ patch("awg_gus.manage.Awg.call", new_callable=AsyncMock): with self.assertRaisesRegex(ValueError, "pairing failed"): manage.setup_profile("default", reconfigure=True, layout=self.layout, setup_fn=setup, check_fn=AsyncMock()) self.assertEqual(self.layout.config("default").read_bytes(), original) s = State(database, "original") self.assertEqual(s.get("offset"), 91) s.db.close() self.assertIn(Layout.unit("default"), systemd.active) def test_check_continues_when_one_gus_is_unreachable(self): save_config(self.layout.config("gus2"), config(self.layout, "gus2", BOT_B)) seen = [] async def fake_check(args): seen.append(args.config) if args.config == str(self.layout.config("default")): raise OSError("offline") with patch("awg_gus.manage.check", fake_check), contextlib.redirect_stdout(io.StringIO()): with self.assertRaisesRegex(ValueError, "default"): manage.dispatch("check", None, self.layout) self.assertEqual(len(seen), 2) def test_failed_database_snapshot_does_not_restore_partial_backup(self): database = self.layout.state("default") / "state.sqlite3" s = State(database, "original") s.set("offset", 91) s.db.close() systemd = FakeSystemd() def bad_copy(source, destination): Path(destination).write_bytes(b"partial backup") raise OSError("disk full") with patch("awg_gus.manage.systemctl", systemd), \ patch("awg_gus.manage.Awg.call", new_callable=AsyncMock), \ patch("awg_gus.manage.database_copy", side_effect=bad_copy) as copy: with self.assertRaises(OSError): manage.setup_profile("default", reconfigure=True, layout=self.layout) self.assertEqual(copy.call_count, 1) s = State(database, "original") self.assertEqual(s.get("offset"), 91) s.db.close() self.assertIn(Layout.unit("default"), systemd.active) class ParallelBotsTests(unittest.IsolatedAsyncioTestCase): async def test_parallel_gus_updates_confirmations_and_owners_are_isolated(self): with tempfile.TemporaryDirectory() as directory: states = [State(Path(directory)/f"{i}.sqlite3", f"gus{i}") for i in range(2)] api = FakeAwg() deliveries = [FakeDelivery(), FakeDelivery()] bots = [Bot({"owners": [owner]}, api, delivery, state) for owner, delivery, state in zip([OWNER, OTHER], deliveries, states)] async def feed(bot, event): await bot.process(event) await asyncio.gather(*bot.tasks.values()) try: await asyncio.gather(feed(bots[0], callback(1, "delete:alice")), feed(bots[1], message(1, "/menu", OTHER))) token = next(iter(bots[0].confirmations)) await feed(bots[1], callback(2, "yes:"+token, OTHER)) self.assertNotIn(("client", "del", "alice"), api.calls) await feed(bots[0], callback(2, "yes:"+token)) await feed(bots[1], message(3, "/clients", OTHER)) self.assertEqual(api.rows, []) # mutation visible from both connections before = len(api.calls) await feed(bots[1], callback(4, "conf:alice", OWNER)) self.assertEqual(len(api.calls), before) self.assertEqual(states[0].get("offset"), 3) self.assertEqual(states[1].get("offset"), 5) finally: for state in states: state.db.close() if __name__ == "__main__": unittest.main()