import datetime from unittest import mock from django.test import TestCase from django.utils import timezone from dramatiq.middleware.time_limit import TimeLimitExceeded from Telegram import interfaces from Telegram.models import Bot, Chat, User as TelegramUser from Telegram.testing import Conversation, FakeApi from . import fields, services, tasks from .bot import texts from .bot.interface import MyPointInterface, is_proxy_url, parse_credentials from .models import Credentials, Point, Server from .tests_services import FakeRunner, FakeXuiApi, create_point, patch_runner, patch_xui_api class ParsersTests(TestCase): def test_credentials(self): self.assertEqual(parse_credentials('1.2.3.4:pa:ss').password, 'pa:ss') self.assertIsNone(parse_credentials('1.2.3.4:')) self.assertIsNone(parse_credentials('host:pw')) self.assertIsNone(parse_credentials('1.2.3.400:pw')) def test_proxy_url(self): self.assertTrue(is_proxy_url('http://127.0.0.1:9050')) self.assertTrue(is_proxy_url('socks5://user:pw@proxy.example.com:1080')) self.assertFalse(is_proxy_url('http://127.0.0.1')) self.assertFalse(is_proxy_url('ftp://127.0.0.1:21')) class BotTestCase(TestCase): def setUp(self): self.bot = Bot.objects.create(token='1:test', interface=interfaces.key(MyPointInterface)) self.chat = Conversation(self.bot) self.api = self.chat.api def say(self, text): return self.chat.say(text) @property def keyboard(self): return self.api.last_keyboard class MainMenuTests(BotTestCase): def test_start_and_main_menu_button_are_same(self): self.assertEqual(self.say('/start'), texts.GREETING) self.assertEqual(self.keyboard, texts.MAIN_KEYBOARD) self.assertEqual(self.say(texts.MAIN_MENU), texts.GREETING) def test_lists_points(self): self.say('/start') point = create_point(TelegramUser.objects.get()) self.assertIn(point.name, self.say('/start')) self.assertEqual(self.keyboard, texts.MAIN_KEYBOARD_WITH_POINTS) def test_faq_and_help(self): self.assertEqual(self.say(texts.FAQ), texts.FAQ_TEXT) self.assertEqual(self.say(texts.HELP), texts.HELP_TEXT) @mock.patch.object(tasks.install_point, 'send') class AddServerTests(BotTestCase): def fill(self, role_button, credentials): self.say(role_button) return self.say(credentials) def test_full_flow(self, send): self.say(texts.ADD_SERVER) self.assertEqual(self.keyboard, texts.DRAFT_KEYBOARD) card = self.fill(texts.SERVER, '10.0.0.1:panel-pw') self.assertIn('Сервер: 10.0.0.1', card) self.assertNotIn('panel-pw', card) self.fill(texts.MASK, '10.0.0.2:mask-pw') self.say(texts.DOMAIN) self.say('Example.COM.') self.say(texts.PROXY) card = self.say(texts.CF_WARP) self.assertIn('Домен: example.com', card) self.assertIn('Прокси: cf_warp', card) self.assertEqual(self.say(texts.INSTALL), texts.INSTALL_STARTED) point = Point.objects.get() self.assertEqual((point.panel.address, point.panel.root_password), ('10.0.0.1', 'panel-pw')) self.assertEqual((point.mask.address, point.mask.root_password), ('10.0.0.2', 'mask-pw')) self.assertEqual((point.domain, point.outbound_proxy), ('example.com', 'cf_warp')) self.assertEqual(point.owner, TelegramUser.objects.get()) send.assert_called_once_with(point.pk, mock.ANY) point.refresh_from_db() self.assertEqual((point.task_chat, point.task_token), (Chat.objects.get(), send.call_args.args[1])) # сообщения с паролями удалены из чата self.assertEqual(len(self.api.deleted), 2) self.assertEqual(Chat.objects.get().state, '') def test_validation(self, send): self.say(texts.ADD_SERVER) self.assertEqual(self.say(texts.INSTALL), texts.DRAFT_INCOMPLETE) self.assertEqual(self.fill(texts.SERVER, 'nonsense'), texts.CREDENTIALS_INVALID) self.say('10.0.0.1:pw') self.assertEqual(self.fill(texts.MASK, '10.0.0.1:pw'), texts.ADDRESS_DUPLICATE) create_point(TelegramUser.objects.get()) self.assertEqual(self.say('10.0.0.2:pw'), texts.ADDRESS_TAKEN.format(address='10.0.0.2')) self.say(texts.CANCEL) self.assertEqual(self.fill(texts.DOMAIN, 'not a domain'), texts.DOMAIN_INVALID) self.say(texts.CANCEL) self.assertEqual(self.fill(texts.PROXY, 'http://nope'), texts.PROXY_INVALID) send.assert_not_called() def test_install_without_mask(self, send): self.say(texts.ADD_SERVER) self.fill(texts.SERVER, '10.0.0.1:pw') self.assertEqual(self.say(texts.INSTALL), texts.INSTALL_STARTED) point = Point.objects.get() self.assertIsNone(point.mask) send.assert_called_once_with(point.pk, mock.ANY) def test_mask_requires_domain(self, send): self.say(texts.ADD_SERVER) self.fill(texts.SERVER, '10.0.0.1:pw') self.fill(texts.MASK, '10.0.0.2:pw') self.assertEqual(self.say(texts.INSTALL), texts.MASK_WITHOUT_DOMAIN) self.say(texts.MASK) self.assertIn('Сервер маскировки: -', self.say(texts.DELETE)) self.assertEqual(self.say(texts.INSTALL), texts.INSTALL_STARTED) self.assertIsNone(Point.objects.get().mask) def test_cancel_in_prompt_returns_to_draft(self, send): self.say(texts.ADD_SERVER) self.fill(texts.SERVER, '10.0.0.1:pw') self.say(texts.MASK) self.assertIn('Сервер: 10.0.0.1', self.say(texts.CANCEL)) self.assertEqual(self.say(texts.CANCEL), texts.GREETING) class PointMenuTests(BotTestCase): def setUp(self): super().setUp() self.say('/start') self.point = create_point(TelegramUser.objects.get()) self.point.set_status(Point.Status.READY) def open_point(self): self.say(texts.MANAGE) return self.say(self.point.name) def test_select_point(self): self.assertIn(self.point.name, self.say(texts.MANAGE)) self.assertEqual(self.say('unknown'), texts.POINT_NOT_FOUND) card = self.say(self.point.name) self.assertIn('Состояние: активен', card) self.assertEqual(self.keyboard, texts.POINT_KEYBOARD) def test_foreign_point_is_not_found(self): other = Conversation(self.bot, user_id=8) other.say(texts.MANAGE) self.assertEqual(other.say(self.point.name), texts.POINT_NOT_FOUND) @mock.patch.object(tasks.release_point, 'send') def test_release(self, send): self.open_point() self.say(texts.RELEASE) self.assertIn(self.point.name, self.say(texts.NO)) self.assertIn('Продолжить?', self.say(texts.RELEASE)) self.assertEqual(self.say(texts.YES), texts.RELEASE_STARTED) send.assert_called_once_with(self.point.pk, mock.ANY) self.point.refresh_from_db() self.assertEqual(self.point.status, Point.Status.RELEASING) @mock.patch.object(tasks.release_point, 'send') def test_busy_point(self, send): self.point.set_status(Point.Status.INSTALLING) self.open_point() self.assertEqual(self.say(texts.RELEASE), texts.POINT_BUSY.format(status='устанавливается')) send.assert_not_called() @mock.patch.object(tasks.reinstall_point, 'send') def test_reinstall_one_server(self, send): self.open_point() self.say(texts.REINSTALL) self.assertEqual(self.say('10.9.9.9:pw'), texts.REINSTALL_UNKNOWN_ADDRESS.format(address='10.9.9.9')) self.assertEqual(self.say('10.0.0.2:new-pw'), texts.REINSTALL_STARTED) send.assert_called_once_with(self.point.pk, mock.ANY, '10.0.0.2', 'new-pw') self.assertEqual(len(self.api.deleted), 2) @mock.patch.object(tasks.reinstall_point, 'send') def test_reinstall_with_old_passwords(self, send): self.open_point() self.say(texts.REINSTALL) self.say(texts.USE_OLD) send.assert_called_once_with(self.point.pk, mock.ANY, None, None) def test_server_menu_button_opens_latest_point(self): self.assertIn(self.point.name, self.say(texts.SERVER_MENU)) class KeysMenuTests(BotTestCase): def setUp(self): super().setUp() self.say('/start') self.point = create_point(TelegramUser.objects.get(), mask=None) self.point.panel_inbound_id = 1 self.point.status = Point.Status.READY self.point.save() self.panel_api = FakeXuiApi() patcher = patch_xui_api(self.panel_api) patcher.start() self.addCleanup(patcher.stop) self.say(texts.MANAGE) self.say(self.point.name) def test_issue_show_and_revoke_key(self): self.assertIn('Ключей пока нет', self.say(texts.CLIENTS)) self.assertEqual(self.keyboard, texts.KEYS_KEYBOARD) self.assertEqual(self.say(texts.ADD_KEY), texts.KEY_NAME_PROMPT) self.assertIn('латиница', self.say('Мама')) card = self.say('Mom-Phone') self.assertIn(self.point.subscription_url('sub-mom-phone'), card) self.assertIn('mom-phone', self.say(texts.CLIENTS)) self.assertIn('↓ 3.0 МБ', self.say('mom-phone')) self.assertEqual(self.say('ghost'), texts.KEY_NOT_FOUND) self.say(texts.REVOKE_KEY) self.assertIn('Ключей пока нет', self.say('mom-phone')) self.assertEqual(self.panel_api.clients, {}) def test_panel_error(self): self.say(texts.CLIENTS) self.say(texts.ADD_KEY) self.panel_api.down = True self.assertEqual(self.say('mom'), texts.PANEL_ERROR.format(error='connection refused')) def test_not_ready_point(self): self.point.set_status(Point.Status.ERROR) self.assertEqual(self.say(texts.CLIENTS), texts.POINT_NOT_READY) def test_card_counts_keys(self): self.say(texts.CLIENTS) self.say(texts.ADD_KEY) self.say('mom') self.assertIn('Ключи: 1', self.say(texts.SERVER_MENU)) class PointSettingsTests(BotTestCase): def setUp(self): super().setUp() self.say('/start') self.point = create_point(TelegramUser.objects.get(), mask=None) self.point.panel_inbound_id = 1 self.point.status = Point.Status.READY self.point.save() self.say(texts.MANAGE) self.say(self.point.name) @mock.patch.object(tasks.set_domain, 'send') def test_change_domain(self, send): self.assertIn('10.0.0.1', self.say(texts.POINT_DOMAIN)) self.assertEqual(self.keyboard, [[texts.CANCEL, texts.DELETE]]) self.assertEqual(self.say('not a domain'), texts.DOMAIN_INVALID) self.assertEqual(self.say('VPN.Example.com.'), texts.CHANGE_STARTED) send.assert_called_once_with(self.point.pk, mock.ANY, 'vpn.example.com') self.point.refresh_from_db() self.assertEqual(self.point.status, Point.Status.UPDATING) self.assertEqual(self.say(texts.SERVER_MENU).count('настраивается'), 1) self.assertEqual(self.say(texts.POINT_DOMAIN), texts.POINT_BUSY.format(status='настраивается')) @mock.patch.object(tasks.set_domain, 'send') def test_mask_point_cannot_drop_domain(self, send): Server.objects.create_server(self.point, Server.Role.MASK, Credentials('10.0.0.3', 'pw')) self.point.domain = 'vpn.example.com' self.point.save() self.assertIn('10.0.0.3', self.say(texts.POINT_DOMAIN)) self.assertEqual(self.say(texts.DELETE), texts.DOMAIN_REQUIRED_BY_MASK) send.assert_not_called() @mock.patch.object(tasks.set_domain, 'send') def test_drop_domain(self, send): self.point.domain = 'vpn.example.com' self.point.save() self.say(texts.POINT_DOMAIN) self.say(texts.DELETE) send.assert_called_once_with(self.point.pk, mock.ANY, '') @mock.patch.object(tasks.set_outbound_proxy, 'send') def test_change_proxy(self, send): self.say(texts.POINT_PROXY) self.assertEqual(self.keyboard, [[texts.CF_WARP, texts.CANCEL, texts.DELETE]]) self.assertEqual(self.say('ftp://x:21'), texts.PROXY_INVALID) self.assertEqual(self.say('socks5://user:secret@proxy.example:1080'), texts.CHANGE_STARTED) # сообщение с паролем удалено, в брокер прокси уходит зашифрованным self.assertEqual(len(self.api.deleted), 2) encrypted = send.call_args.args[2] self.assertNotIn('secret', encrypted) self.assertEqual(fields.decrypt(encrypted), 'socks5://user:secret@proxy.example:1080') @mock.patch.object(tasks.set_outbound_proxy, 'send') def test_warp_and_direct(self, send): self.say(texts.POINT_PROXY) self.say(texts.CF_WARP) self.assertEqual(fields.decrypt(send.call_args.args[2]), Point.CF_WARP) self.point.set_status(Point.Status.READY) self.say(texts.SERVER_MENU) self.say(texts.POINT_PROXY) self.say(texts.DELETE) self.assertEqual(send.call_args.args[2], '') def test_cancel_returns_to_point(self): self.say(texts.POINT_PROXY) self.assertIn('Состояние: активен', self.say(texts.CANCEL)) self.assertEqual(self.keyboard, texts.POINT_KEYBOARD) def test_not_ready_point(self): self.point.set_status(Point.Status.ERROR) self.assertEqual(self.say(texts.POINT_PROXY), texts.POINT_NOT_READY) def test_card_hides_proxy_password(self): self.point.outbound_proxy = 'socks5://user:secret@proxy.example:1080' self.point.save() card = self.say(texts.SERVER_MENU) self.assertIn('socks5://***@proxy.example:1080', card) self.assertNotIn('secret', card) self.assertNotIn('user', card) class TasksTests(TestCase): def setUp(self): bot = Bot.objects.create(token='1:test') owner = TelegramUser.objects.create(tg_id=1) self.chat = Chat.objects.create(bot=bot, chat_id=1, type='private', user=owner) self.point = create_point(owner) self.api = FakeApi() def run_task(self, actor, *args, status=Point.Status.INSTALLING, runner=None): """Запускает задачу, как бот, и выполняет ее сообщение.""" self.point.refresh_from_db() with mock.patch.object(actor, 'send') as send: tasks.start(actor, self.point, status, self.chat, *args) with ( patch_runner(runner or FakeRunner()), mock.patch.object(Bot, 'api', return_value=self.api), mock.patch.object(tasks.check_point, 'send'), ): actor.fn(*send.call_args.args) return self.api.last_text def test_install_notifies_success(self): text = self.run_task(tasks.install_point) self.assertIn(self.point.name, text) self.assertIn('10.0.0.2', text) self.assertIn('готов к работе', text) self.assertEqual(self.api.last_keyboard, [[texts.SERVER_MENU]]) def test_install_notifies_failure(self): with self.assertLogs('App.tasks', 'ERROR'): text = self.run_task(tasks.install_point, runner=FakeRunner(fail='xui')) self.assertIn('Не удалось', text) self.assertIn('boom', text) def test_release_notifies_passwords(self): text = self.run_task(tasks.release_point, status=Point.Status.RELEASING) self.assertIn('10.0.0.1 - пароль root', text) self.assertFalse(Point.objects.exists()) def test_change_notifies_result(self): self.point.panel_inbound_id = 1 self.point.status = Point.Status.READY self.point.save() proxy = fields.encrypt('socks5://user:secret@proxy.example:1080') text = self.run_task(tasks.set_outbound_proxy, proxy, status=Point.Status.UPDATING) self.assertIn('socks5://***@proxy.example', text) with self.assertLogs('App.tasks', 'ERROR'): text = self.run_task( tasks.set_domain, 'vpn.example.com', status=Point.Status.UPDATING, runner=FakeRunner(fail='xui_cert'), ) self.assertIn('прежними настройками', text) self.point.refresh_from_db() self.assertEqual((self.point.status, self.point.domain), (Point.Status.READY, '')) class SupervisorTests(TestCase): """Задача, которая прервалась, не оставляет точку занятой навсегда.""" def setUp(self): bot = Bot.objects.create(token='1:test') owner = TelegramUser.objects.create(tg_id=1) self.chat = Chat.objects.create(bot=bot, chat_id=1, type='private', user=owner) self.point = create_point(owner) self.api = FakeApi() patcher = mock.patch.object(Bot, 'api', return_value=self.api) patcher.start() self.addCleanup(patcher.stop) def start(self, actor=tasks.install_point, status=Point.Status.INSTALLING, *args): with mock.patch.object(actor, 'send') as send, mock.patch.object(tasks.supervise, 'send_with_options') as sup: tasks.start(actor, self.point, status, self.chat, *args) sup.assert_called_once_with(args=('App.Point', self.point.pk, send.call_args.args[1]), delay=tasks.SUPERVISE_DELAY_MS) return send.call_args.args def supervise(self, token): with mock.patch.object(tasks.supervise, 'send_with_options') as again: tasks.supervise.fn('App.Point', self.point.pk, token) self.point.refresh_from_db() return again.called def age_heartbeat(self, delta): Point.objects.filter(pk=self.point.pk).update(task_heartbeat=timezone.now() - delta) def test_stuck_task_is_aborted(self): _, token = self.start() Point.objects.claim_task(self.point.pk, token) self.assertTrue(self.supervise(token)) self.age_heartbeat(tasks.TASK_STALE + datetime.timedelta(seconds=1)) self.assertFalse(self.supervise(token)) self.assertEqual(self.point.status, Point.Status.ERROR) self.assertIn('Переустановить', self.point.error) self.assertEqual(self.point.task_token, '') self.assertIn('установка прервалась', self.api.last_text) def test_queued_task_waits_longer(self): _, token = self.start() self.age_heartbeat(tasks.TASK_STALE * 2) self.assertTrue(self.supervise(token)) self.age_heartbeat(tasks.QUEUE_TIMEOUT + datetime.timedelta(seconds=1)) self.assertFalse(self.supervise(token)) self.assertEqual(self.point.status, Point.Status.ERROR) # сообщение все же дошло: задача снята, работы нет runner = FakeRunner() with patch_runner(runner): tasks.install_point.fn(self.point.pk, token) self.assertEqual(runner.calls, []) def test_interrupted_update_keeps_point_ready(self): token = self.start(tasks.set_domain, Point.Status.UPDATING, 'vpn.example.com')[1] Point.objects.claim_task(self.point.pk, token) self.age_heartbeat(tasks.TASK_STALE * 2) self.supervise(token) self.assertEqual((self.point.status, self.point.domain), (Point.Status.READY, '')) self.assertIn('прервалась', self.point.error) def test_supervisor_stops_after_task(self): args = self.start() with patch_runner(FakeRunner()), mock.patch.object(tasks.check_point, 'send'): tasks.install_point.fn(*args) self.assertFalse(self.supervise(args[1])) self.assertEqual((self.point.status, self.point.task_token), (Point.Status.READY, '')) def test_redelivered_message_is_not_repeated(self): # воркер упал посреди установки, RabbitMQ доставил сообщение снова args = self.start() Point.objects.claim_task(*args) runner = FakeRunner() with patch_runner(runner): tasks.install_point.fn(*args) self.assertEqual(runner.calls, []) self.point.refresh_from_db() self.assertEqual(self.point.status, Point.Status.ERROR) def test_time_limit_aborts_task(self): args = self.start() with ( mock.patch.object(services, 'install_point', side_effect=TimeLimitExceeded), self.assertRaises(TimeLimitExceeded), self.assertLogs('App.tasks', 'ERROR'), ): tasks.install_point.fn(*args) self.point.refresh_from_db() self.assertEqual((self.point.status, self.point.task_token), (Point.Status.ERROR, '')) def test_failed_notification_keeps_result(self): args = self.start() with ( patch_runner(FakeRunner()), mock.patch.object(tasks, 'schedule_checks', side_effect=RuntimeError('boom')), self.assertRaises(RuntimeError), self.assertLogs('App.tasks', 'ERROR'), ): tasks.install_point.fn(*args) self.point.refresh_from_db() self.assertEqual((self.point.status, self.point.task_token), (Point.Status.READY, '')) class CheckTests(TestCase): def setUp(self): bot = Bot.objects.create(token='1:test') owner = TelegramUser.objects.create(tg_id=1) Chat.objects.create(bot=bot, chat_id=1, type='private', user=owner) self.point = create_point(owner) self.point.status = Point.Status.READY self.point.save() self.api = FakeApi() patcher = mock.patch.object(Bot, 'api', return_value=self.api) patcher.start() self.addCleanup(patcher.stop) with mock.patch.object(tasks.check_point, 'send') as send: tasks.schedule_checks(self.point) self.token = send.call_args.args[1] def check(self, problem='', token=None): with ( mock.patch.object(services, 'check_point', return_value=problem) as check, mock.patch.object(tasks.check_point, 'send_with_options') as again, ): tasks.check_point.fn(self.point.pk, token or self.token) self.point.refresh_from_db() return check.called, again.called def test_checks_once_a_day_and_reports_changes(self): self.assertEqual(self.check('панель не отвечает'), (True, True)) self.assertEqual(self.point.health_error, 'панель не отвечает') self.assertIn('панель не отвечает', self.api.last_text) self.assertIn('Проверка: панель не отвечает', texts.point_card(self.point)) # до следующих суток не проверяет, но цепочка живет self.assertEqual(self.check(), (False, True)) Point.objects.filter(pk=self.point.pk).update(checked_at=timezone.now() - tasks.CHECK_INTERVAL) sent = len(self.api.sent) self.check('панель не отвечает') self.assertEqual(len(self.api.sent), sent) Point.objects.filter(pk=self.point.pk).update(checked_at=timezone.now() - tasks.CHECK_INTERVAL) self.check() self.assertIn('снова все в порядке', self.api.last_text) def test_busy_point_is_not_checked(self): self.point.set_status(Point.Status.INSTALLING) self.assertEqual(self.check(), (False, True)) def test_old_chain_stops(self): with mock.patch.object(tasks.check_point, 'send'): tasks.schedule_checks(self.point) self.assertEqual(self.check(token=self.token), (False, False)) point_id = self.point.pk self.point.delete() with mock.patch.object(tasks.check_point, 'send_with_options') as again: tasks.check_point.fn(point_id, self.token) again.assert_not_called()