241 lines
11 KiB
Python
241 lines
11 KiB
Python
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()
|