zapretgui/tests/test_startup_address_refresh.py

from __future__ import annotations

import unittest
from unittest.mock import patch


def _managed_block(*rows: str) -> str:
    return "\n".join(
        [
            "127.0.0.1 localhost",
            "# >>> zapretgui:hosts managed begin >>>",
            "# Generated by ZapretGUI. Do not edit this block manually.",
            *rows,
            "# <<< zapretgui:hosts managed end <<<",
            "",
        ]
    )


class DnsAddressMigrationPlanTests(unittest.TestCase):
    def setUp(self) -> None:
        from dns.dns_providers import OUTDATED_DNS_ADDRESS_REPLACEMENTS

        self.replacements = OUTDATED_DNS_ADDRESS_REPLACEMENTS

    def test_old_xbox_addresses_are_replaced_in_place(self) -> None:
        from dns.address_migration import plan_dns_server_migration

        # Xbox DNS закрыл сервисы ИИ: все его адреса меняются на DNS-AI.
        for old in (
            ["176.99.11.77", "80.78.247.254"],
            ["111.88.96.50", "111.88.96.51"],
            ["111.88.96.54", "111.88.96.55"],
            ["87.228.47.200", "87.228.47.201"],
        ):
            with self.subTest(old=old):
                self.assertEqual(
                    plan_dns_server_migration(old, self.replacements),
                    ["192.144.59.14", "186.246.49.127"],
                )

    def test_foreign_addresses_and_order_are_kept(self) -> None:
        from dns.address_migration import plan_dns_server_migration

        self.assertEqual(
            plan_dns_server_migration(["1.1.1.1", "176.99.11.77", "8.8.8.8"], self.replacements),
            ["1.1.1.1", "192.144.59.14", "8.8.8.8"],
        )

    def test_nothing_to_change_returns_none(self) -> None:
        from dns.address_migration import plan_dns_server_migration

        self.assertIsNone(plan_dns_server_migration([], self.replacements))
        self.assertIsNone(plan_dns_server_migration(["1.1.1.1", "8.8.8.8"], self.replacements))
        self.assertIsNone(plan_dns_server_migration(["192.144.59.14", "186.246.49.127"], self.replacements))

    def test_duplicate_after_replacement_is_removed(self) -> None:
        from dns.address_migration import plan_dns_server_migration

        self.assertEqual(
            plan_dns_server_migration(["111.88.96.50", "111.88.96.54"], self.replacements),
            ["192.144.59.14"],
        )

    def test_ipv6_is_matched_in_any_spelling(self) -> None:
        from dns.address_migration import plan_dns_server_migration

        # В списке замен сейчас только IPv4, поэтому пара IPv6 задана прямо здесь.
        replacements = {"2a12:bec4:1460:d5::2": "2a01:4f9:c014:6dac::1"}
        self.assertEqual(
            plan_dns_server_migration(["2A12:BEC4:1460:D5:0:0:0:2"], replacements),
            ["2a01:4f9:c014:6dac::1"],
        )

    def test_old_addresses_are_not_offered_by_provider_list(self) -> None:
        from dns.dns_providers import DNS_PROVIDERS, OUTDATED_DNS_ADDRESS_REPLACEMENTS

        offered = {
            address
            for group in DNS_PROVIDERS.values()
            for provider in group.values()
            for family in ("ipv4", "ipv6")
            for address in provider.get(family, [])
        }
        for old, new in OUTDATED_DNS_ADDRESS_REPLACEMENTS.items():
            self.assertNotIn(old, offered)
            self.assertIn(new, offered)


class HostsAppliedSelectionRefreshDecisionTests(unittest.TestCase):
    def _decide(self, **kwargs):
        from hosts.hosts import decide_applied_selection_refresh

        params = {
            "block_rows": [("chatgpt.com", "95.182.120.241")],
            "has_saved_selection": True,
            "ipv6_available": True,
            "build_desired_rows": lambda: [("chatgpt.com", "45.88.174.254")],
        }
        params.update(kwargs)
        return decide_applied_selection_refresh(**params)[0]

    def test_outdated_block_is_rewritten(self) -> None:
        self.assertTrue(self._decide())

    def test_without_managed_block_nothing_is_written(self) -> None:
        self.assertFalse(self._decide(block_rows=[]))

    def test_without_saved_selection_nothing_is_written(self) -> None:
        self.assertFalse(self._decide(has_saved_selection=False))

    def test_up_to_date_block_is_not_rewritten(self) -> None:
        self.assertFalse(self._decide(build_desired_rows=lambda: [("ChatGPT.com", "95.182.120.241")]))

    def test_same_addresses_in_another_order_of_names_are_not_rewritten(self) -> None:
        # Страница пишет сервисы в порядке нажатий, сверка строит их в сохранённом порядке.
        block = [("x.com", "1.1.1.1"), ("x.com", "1.1.1.2"), ("github.com", "2.2.2.2")]
        self.assertFalse(
            self._decide(
                block_rows=block,
                build_desired_rows=lambda: [("github.com", "2.2.2.2"), ("x.com", "1.1.1.1"), ("x.com", "1.1.1.2")],
            )
        )

    def test_another_first_address_of_a_name_is_a_real_change(self) -> None:
        # Браузер берёт первую строку имени: смена первого адреса — повод переписать.
        block = [("x.com", "1.1.1.1"), ("x.com", "1.1.1.2")]
        self.assertTrue(self._decide(block_rows=block, build_desired_rows=lambda: [("x.com", "1.1.1.2"), ("x.com", "1.1.1.1")]))

    def test_ipv6_rows_are_not_dropped_while_ipv6_is_down(self) -> None:
        self.assertFalse(
            self._decide(
                block_rows=[("github.com", "2606:50c0:8000::154")],
                ipv6_available=False,
            )
        )

    def test_selection_that_became_empty_clears_dead_rows(self) -> None:
        self.assertTrue(self._decide(build_desired_rows=lambda: []))


class HostsAppliedSelectionRefreshRunTests(unittest.TestCase):
    def test_outdated_block_is_rewritten_with_catalog_addresses(self) -> None:
        from hosts import hosts as hosts_module

        written: list[str] = []
        manager = hosts_module.HostsManager(status_callback=lambda _message: None)
        manager.is_hosts_file_accessible = lambda: True
        original = _managed_block("95.182.120.241 chatgpt.com")

        with (
            patch.object(hosts_module, "safe_read_hosts_file", return_value=original),
            patch.object(hosts_module, "safe_write_hosts_file", side_effect=lambda text: written.append(text) or True),
            patch.object(hosts_module, "is_ipv6_available", return_value=True),
            patch.object(
                hosts_module,
                "get_service_domain_ip_rows",
                return_value=[("chatgpt.com", "45.88.174.254")],
            ),
        ):
            changed, _reason = manager.refresh_applied_service_selection(
                {"ChatGPT & Sora (OpenAI)": "comss_dns"},
                has_saved_selection=True,
            )

        self.assertTrue(changed)
        self.assertEqual(len(written), 1)
        self.assertIn("45.88.174.254 chatgpt.com", written[0])
        self.assertNotIn("95.182.120.241", written[0])
        self.assertIn("127.0.0.1 localhost", written[0])

    def test_up_to_date_block_is_left_untouched(self) -> None:
        from hosts import hosts as hosts_module

        manager = hosts_module.HostsManager(status_callback=lambda _message: None)
        manager.is_hosts_file_accessible = lambda: (_ for _ in ()).throw(
            AssertionError("актуальный hosts не должен проверяться на запись")
        )
        original = _managed_block("45.88.174.254 chatgpt.com")

        with (
            patch.object(hosts_module, "safe_read_hosts_file", return_value=original),
            patch.object(hosts_module, "safe_write_hosts_file") as write_hosts,
            patch.object(hosts_module, "is_ipv6_available", return_value=True),
            patch.object(
                hosts_module,
                "get_service_domain_ip_rows",
                return_value=[("chatgpt.com", "45.88.174.254")],
            ),
        ):
            changed, _reason = manager.refresh_applied_service_selection(
                {"ChatGPT & Sora (OpenAI)": "comss_dns"},
                has_saved_selection=True,
            )

        self.assertFalse(changed)
        write_hosts.assert_not_called()

    def test_command_reports_change_and_uses_saved_selection(self) -> None:
        from hosts import commands as hosts_commands

        calls: list[tuple[dict[str, str], bool]] = []

        class FakeManager:
            def refresh_applied_service_selection(self, service_dns, *, has_saved_selection):
                calls.append((dict(service_dns), has_saved_selection))
                return True, "адреса в hosts устарели"

        with (
            patch.object(hosts_commands, "load_user_selection", return_value={"Claude": "xbox_dns"}),
            patch("hosts.proxy_domains.has_saved_user_hosts_selection", return_value=True),
        ):
            result = hosts_commands.refresh_applied_selection(FakeManager())

        self.assertTrue(result.changed)
        self.assertEqual(calls, [({"Claude": "xbox_dns"}, True)])


class StartupAddressRefreshInstallTests(unittest.TestCase):
    def _run_install(self, module_name: str, install_name: str, feature_kwargs: dict, *, gate_allows: bool = True) -> list[tuple[str, str]]:
        import importlib

        from utils import probe_gate

        module = importlib.import_module(module_name)
        queued: list[tuple[str, str, object]] = []

        class _State:
            interactive_logged = True

        class _Host:
            startup_state = _State()
            startup_interactive_ready = object()

        with (
            patch.object(module, "enqueue_subsystem_task", side_effect=lambda q, n, t: queued.append((q, n, t))),
            patch.object(module, "bind_startup_gate", side_effect=lambda _signal, cb, is_ready: cb()),
            patch.object(module, "is_startup_host_alive", return_value=True),
            # Повтор по расписанию ставится таймером Qt: в тесте его нет, запоминаем только срок.
            patch.object(module, "schedule_after", side_effect=lambda delay, _cb: self.scheduled.append(delay), create=True),
            # Шлюз проб: разрешает плановую проверку или нет (время прошлой — в настройках).
            patch.object(
                probe_gate,
                "try_background",
                side_effect=lambda kind: probe_gate.BackgroundPermit(kind, probe_gate.Why.SCHEDULE) if gate_allows else None,
            ),
            patch.object(probe_gate.BackgroundPermit, "release", lambda self: None),
        ):
            getattr(module, install_name)(_Host(), log_startup_metric=lambda *_args: None, **feature_kwargs)
            # Задача может поставить следующую: проверка адресов ставит запись в hosts.
            index = 0
            while index < len(queued):
                queued[index][2]()
                index += 1
        return [(queue, name) for queue, name, _target in queued]

    def setUp(self) -> None:
        self.scheduled: list[int] = []

    def test_addresses_are_checked_in_a_waiting_queue_and_only_then_hosts_is_rewritten(self) -> None:
        calls: list[str] = []

        from utils.probe_gate import Why

        class FakeHostsFeature:
            def check_addresses(self, *, why):
                calls.append(f"check:{why}")

            def refresh_applied_selection(self):
                calls.append("refresh")

        queued = self._run_install(
            "main.post_startup_hosts_refresh",
            "install_hosts_applied_selection_refresh",
            {"hosts_feature": FakeHostsFeature()},
        )
        # Проверка ждёт сеть — в своей очереди; запись файла — на общей дорожке, после проверки.
        self.assertEqual(queued, [("dns", "HostsAddressCheck"), ("hosts", "HostsAppliedSelectionRefresh")])
        self.assertEqual(calls, [f"check:{Why.SCHEDULE}", "refresh"])
        # Дальше раз в час спрашиваем шлюз проб; сам интервал (6 ч) решает он.
        self.assertEqual(self.scheduled, [3600 * 1000])

    def test_gate_refusal_skips_the_address_check_but_not_the_hosts_refresh(self) -> None:
        # Перезапуск через час после прошлой проверки: шлюз не пускает — ни одной пробы.
        calls: list[str] = []

        class FakeHostsFeature:
            def check_addresses(self, *, why):
                calls.append("check")

            def refresh_applied_selection(self):
                calls.append("refresh")

        self._run_install(
            "main.post_startup_hosts_refresh",
            "install_hosts_applied_selection_refresh",
            {"hosts_feature": FakeHostsFeature()},
            gate_allows=False,
        )
        self.assertEqual(calls, ["refresh"])

    def test_failed_address_check_does_not_cancel_the_hosts_refresh(self) -> None:
        calls: list[str] = []

        class FakeHostsFeature:
            def check_addresses(self, *, why):
                raise OSError("нет сети")

            def refresh_applied_selection(self):
                calls.append("refresh")

        self._run_install(
            "main.post_startup_hosts_refresh",
            "install_hosts_applied_selection_refresh",
            {"hosts_feature": FakeHostsFeature()},
        )
        self.assertEqual(calls, ["refresh"])

    def test_hosts_refresh_logs_why_nothing_changed(self) -> None:
        import importlib
        from types import SimpleNamespace
        from unittest.mock import patch

        module = importlib.import_module("main.post_startup_hosts_refresh")
        metrics: list[tuple[str, str]] = []
        feature = SimpleNamespace(
            refresh_applied_selection=lambda: SimpleNamespace(changed=False, message="выбор сервисов не сохранён")
        )
        host = SimpleNamespace(startup_interactive_ready=None, startup_state=SimpleNamespace(interactive_logged=True))
        with (
            patch.object(module, "enqueue_subsystem_task", side_effect=lambda _queue, _name, target: target()),
            patch.object(module, "bind_startup_gate", side_effect=lambda _signal, cb, is_ready: cb()),
            patch.object(module, "is_startup_host_alive", return_value=True),
        ):
            module.install_hosts_applied_selection_refresh(
                host, hosts_feature=feature, log_startup_metric=lambda name, text: metrics.append((name, text))
            )
        self.assertEqual(
            metrics, [("StartupHostsAppliedSelectionRefreshFinished", "changed=False (выбор сервисов не сохранён)")]
        )

    def test_dns_migration_runs_in_dns_queue(self) -> None:
        calls: list[str] = []

        class FakeDnsFeature:
            def migrate_outdated_dns_addresses(self):
                calls.append("migrate")
                return []

        queued = self._run_install(
            "main.post_startup_dns_migration",
            "install_dns_address_migration",
            {"dns_feature": FakeDnsFeature()},
        )
        self.assertEqual(queued, [("dns", "DnsAddressMigration")])
        self.assertEqual(calls, ["migrate"])

    def test_post_startup_runs_refresh_before_page_warmups(self) -> None:
        from unittest.mock import Mock

        from main import post_startup
        from main.post_startup import PostStartupDeps, install_post_startup_tasks

        startup_host = object()
        hosts_feature = object()
        dns_feature = object()
        log_startup_metric = Mock()
        deps = PostStartupDeps(
            startup_host=startup_host,
            hosts_feature=hosts_feature,
            profile_feature=object(),
            dns_feature=dns_feature,
            notify=Mock(),
            notify_many=Mock(),
            set_status=Mock(),
            log_startup_metric=log_startup_metric,
            start_proxy_if_enabled_async=Mock(),
            startup_lists_check=Mock(),
            install_tray_post_startup=Mock(),
            updater_feature=Mock(),
        )
        order = Mock()
        installers = [
            name
            for name in dir(post_startup)
            if name.startswith("install_") and name != "install_post_startup_tasks"
        ]
        patches = [patch.object(post_startup, name, getattr(order, name)) for name in installers]
        patches.append(patch.object(post_startup, "build_idle_ui_task_queue", return_value=object()))
        for item in patches:
            item.start()
        try:
            install_post_startup_tasks(deps)
        finally:
            for item in patches:
                item.stop()

        called = [call[0] for call in order.mock_calls]
        self.assertLess(called.index("install_dns_address_migration"), called.index("install_dns_page_data_warmup"))
        self.assertLess(
            called.index("install_hosts_applied_selection_refresh"),
            called.index("install_hosts_page_warmup"),
        )
        order.install_dns_address_migration.assert_called_once_with(
            startup_host,
            dns_feature=dns_feature,
            log_startup_metric=log_startup_metric,
        )
        order.install_hosts_applied_selection_refresh.assert_called_once_with(
            startup_host,
            hosts_feature=hosts_feature,
            log_startup_metric=log_startup_metric,
        )


if __name__ == "__main__":
    unittest.main()