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_malw_and_xbox_addresses_are_replaced_in_place(self) -> None:
        from dns.address_migration import plan_dns_server_migration

        self.assertEqual(
            plan_dns_server_migration(["84.21.189.133", "64.188.98.242"], self.replacements),
            ["95.216.204.218", "80.253.249.40"],
        )
        self.assertEqual(
            plan_dns_server_migration(["111.88.96.50", "111.88.96.51"], self.replacements),
            ["111.88.96.54", "111.88.96.55"],
        )

    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", "84.21.189.133", "8.8.8.8"], self.replacements),
            ["1.1.1.1", "95.216.204.218", "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(["111.88.96.54", "111.88.96.55"], 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(["84.21.189.133", "95.216.204.218"], self.replacements),
            ["95.216.204.218"],
        )

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

        self.assertEqual(
            plan_dns_server_migration(["2A12:BEC4:1460:D5:0:0:0:2"], self.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_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) -> list[tuple[str, str]]:
        import importlib

        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),
        ):
            getattr(module, install_name)(_Host(), log_startup_metric=lambda *_args: None, **feature_kwargs)
            for _queue, _name, target in queued:
                target()
        return [(queue, name) for queue, name, _target in queued]

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

        class FakeHostsFeature:
            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, [("hosts", "HostsAppliedSelectionRefresh")])
        self.assertEqual(calls, ["refresh"])

    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]
        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()