Spaces:
Running
Running
| import asyncio | |
| import dataclasses | |
| import logging | |
| import socket | |
| from types import SimpleNamespace | |
| import aiohttp | |
| import pytest | |
| from aiohttp.test_utils import TestClient, TestServer | |
| from telethon import functions | |
| from telethon.crypto import AuthKey | |
| import bot.main as bot_main | |
| from bot import __version__ | |
| from bot.config import load_config | |
| from bot.db import Database | |
| from bot.errorlog import ChannelLogHandler | |
| from bot.web import make_app, start_web | |
| from .tgfakes import FakeClient | |
| def free_port() -> int: | |
| with socket.socket() as s: | |
| s.bind(("127.0.0.1", 0)) | |
| return s.getsockname()[1] | |
| # --------------------------------------------------------------------------- | |
| # health server | |
| # --------------------------------------------------------------------------- | |
| async def test_health_endpoints_answer_get_and_head(): | |
| async with TestClient(TestServer(make_app())) as client: | |
| r = await client.get("/") | |
| assert r.status == 200 and "running" in await r.text() | |
| assert (await client.head("/")).status == 200 # cron pingers often use HEAD | |
| r = await client.get("/health") | |
| body = await r.json() | |
| assert r.status == 200 and body["status"] == "ok" and body["version"] == __version__ | |
| assert body["uptime_seconds"] >= 0 | |
| assert (await client.get("/nope")).status == 404 | |
| async def test_start_web_listens_on_the_given_port(): | |
| port = free_port() | |
| runner = await start_web(port, host="127.0.0.1") | |
| try: | |
| async with aiohttp.ClientSession() as http: | |
| async with http.get(f"http://127.0.0.1:{port}/health") as r: | |
| assert r.status == 200 | |
| finally: | |
| await runner.cleanup() | |
| # --------------------------------------------------------------------------- | |
| # start-up helpers | |
| # --------------------------------------------------------------------------- | |
| def test_clean_temp_dir_removes_only_job_leftovers(tmp_path): | |
| root = tmp_path / "work" | |
| (root / "job1").mkdir(parents=True) | |
| (root / "job1" / "movie.mkv").write_bytes(b"1") | |
| (root / "tracks5_6").mkdir() | |
| (root / "tracks5_6" / "probe.mkv").write_bytes(b"2") | |
| (root / "other").mkdir() | |
| (root / "keep.txt").write_text("k") | |
| bot_main.clean_temp_dir(str(root)) | |
| assert sorted(p.name for p in root.iterdir()) == ["keep.txt", "other"] | |
| def test_clean_temp_dir_creates_a_missing_directory(tmp_path): | |
| bot_main.clean_temp_dir(str(tmp_path / "a" / "b")) | |
| assert (tmp_path / "a" / "b").is_dir() | |
| def test_main_exits_with_status_2_when_configuration_is_missing(monkeypatch, capsys): | |
| for key in ("API_ID", "API_HASH", "BOT_TOKEN", "OWNER_ID"): | |
| monkeypatch.delenv(key, raising=False) | |
| with pytest.raises(SystemExit) as exc: | |
| bot_main.main() | |
| assert exc.value.code == 2 | |
| assert "Missing required environment variables" in capsys.readouterr().err | |
| class FakeTG: | |
| created: list = [] | |
| def __init__(self, session, api_id, api_hash, **kw): | |
| self.session, self.api_id, self.api_hash, self.kw = session, api_id, api_hash, kw | |
| self.had_key = bool(session.auth_key) | |
| self.started_with = None | |
| FakeTG.created.append(self) | |
| async def start(self, bot_token=None): | |
| self.started_with = bot_token | |
| if not self.session.auth_key: # a fresh login creates an auth key | |
| self.session.set_dc(2, "149.154.167.51", 443) | |
| self.session.auth_key = AuthKey(b"\x07" * 256) | |
| return self | |
| async def login(tmp_path, monkeypatch): | |
| FakeTG.created = [] | |
| monkeypatch.setattr("telethon.TelegramClient", FakeTG) | |
| cfg = load_config({"API_ID": "11", "API_HASH": "hash", "BOT_TOKEN": "1234:abc", "OWNER_ID": "1"}) | |
| db = Database(f"sqlite+aiosqlite:///{tmp_path}/login.db") | |
| await db.init() | |
| yield SimpleNamespace(cfg=cfg, db=db) | |
| await db.close() | |
| async def test_first_login_saves_the_session_and_later_starts_reuse_it(login): | |
| first = await bot_main.build_client(login.cfg, login.db) | |
| assert not first.had_key and first.started_with == "1234:abc" and (first.api_id, first.api_hash) == (11, "hash") | |
| assert first.kw["connection_retries"] == -1 and first.kw["auto_reconnect"] is True | |
| assert await login.db.kv_get("session") and await login.db.kv_get("session_bot") == "1234" | |
| second = await bot_main.build_client(login.cfg, login.db) # e.g. after the Space restarted | |
| assert second.had_key and second.session.dc_id == 2 | |
| async def test_a_saved_session_of_another_bot_is_not_reused(login): | |
| await bot_main.build_client(login.cfg, login.db) | |
| other = dataclasses.replace(login.cfg, bot_token="9999:zzz") | |
| client = await bot_main.build_client(other, login.db) | |
| assert not client.had_key and await login.db.kv_get("session_bot") == "9999" | |
| async def test_session_string_setting_wins_over_the_database(login): | |
| saved = (await bot_main.build_client(login.cfg, login.db)).session.save() | |
| await login.db.kv_delete("session") | |
| await login.db.kv_delete("session_bot") | |
| cfg = dataclasses.replace(login.cfg, session_string=saved) | |
| assert (await bot_main.build_client(cfg, login.db)).had_key | |
| # --------------------------------------------------------------------------- | |
| # whole start-up sequence with a fake Telegram client | |
| # --------------------------------------------------------------------------- | |
| async def test_amain_wires_everything_and_exits_nonzero_if_telegram_disconnects(tmp_path, monkeypatch): | |
| port = free_port() | |
| env = { | |
| "API_ID": "1", "API_HASH": "h", "BOT_TOKEN": "1:t", "OWNER_ID": "1000001", "LOG_CHANNEL": "-1001111111111", | |
| "DATABASE_URL": f"sqlite+aiosqlite:///{tmp_path}/main.db", "TEMP_DIR": str(tmp_path / "work"), | |
| "PORT": str(port), "FAST_TRANSFER": "0", | |
| } | |
| for key, value in env.items(): | |
| monkeypatch.setenv(key, value) | |
| (tmp_path / "work" / "job9").mkdir(parents=True) # leftover of a "crashed" run | |
| gate = asyncio.Event() | |
| client = FakeClient() | |
| client.disconnected = False | |
| async def get_me(): | |
| return SimpleNamespace(username="rename_bot") | |
| async def run_until_disconnected(): | |
| await gate.wait() | |
| async def disconnect(): | |
| client.disconnected = True | |
| client.get_me, client.run_until_disconnected, client.disconnect = get_me, run_until_disconnected, disconnect | |
| async def fake_build_client(cfg, db): | |
| return client | |
| monkeypatch.setattr(bot_main, "build_client", fake_build_client) | |
| task = asyncio.ensure_future(bot_main.amain()) | |
| try: | |
| for _ in range(200): # wait for the "started" notice: everything is wired by then | |
| if any("Bot started" in text for _chat, text, _kw in client.sent): | |
| break | |
| await asyncio.sleep(0.05) | |
| else: | |
| pytest.fail("the bot never finished starting") | |
| async with aiohttp.ClientSession() as http: | |
| async with http.get(f"http://127.0.0.1:{port}/health") as r: | |
| assert r.status == 200 | |
| assert len(client.handlers) > 10 | |
| assert any(isinstance(r, functions.bots.SetBotCommandsRequest) for r in client.requests) | |
| assert not (tmp_path / "work" / "job9").exists() | |
| assert client.sent[0][0] == -1001111111111 | |
| started = client.sent[0][1] # the notice also says what the speed depends on (here: parallel transfer is off) | |
| assert "Bot started" in started and "⚡ cryptg ✅" in started and "parallel off ⚠️" in started | |
| assert any(isinstance(h, ChannelLogHandler) for h in logging.getLogger().handlers) | |
| logging.getLogger("bot.something").error("an error somewhere in the bot") # errors go to the log channel as well | |
| for _ in range(200): | |
| if any("an error somewhere in the bot" in text for _chat, text, _kw in client.sent): | |
| break | |
| await asyncio.sleep(0.05) | |
| else: | |
| pytest.fail("the error never reached the log channel") | |
| gate.set() # Telegram connection lost -> exit non-zero so the container restarts | |
| assert await asyncio.wait_for(task, 10) == 1 | |
| finally: | |
| gate.set() | |
| if not task.done(): | |
| task.cancel() | |
| assert client.disconnected | |
| assert not any(isinstance(h, ChannelLogHandler) for h in logging.getLogger().handlers) # detached again on the way out | |
| with pytest.raises(aiohttp.ClientError): | |
| async with aiohttp.ClientSession() as http: | |
| async with http.get(f"http://127.0.0.1:{port}/health"): | |
| pass | |