dy/backend/tests/test_egress_channels.py

from __future__ import annotations

import asyncio
import os
import sys
import time
import unittest
from pathlib import Path
from unittest.mock import AsyncMock, patch

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.dialects import mysql
from sqlalchemy.schema import CreateTable


BACKEND_DIR = Path(__file__).resolve().parents[1]
if str(BACKEND_DIR) not in sys.path:
    sys.path.insert(0, str(BACKEND_DIR))

from rpa_engine.egress_channels import (
    EgressChannel,
    EgressChannelUnavailable,
    EgressSnapshot,
    LocalAddress,
    discover_egress_channels,
    reset_egress_cache_for_tests,
    resolve_send_channels,
)
from rpa_engine.source_bound_proxy import SourceBoundProxy
from models.db_migrate import migrate_accounts_table
from models.models import Account


class EgressChannelTests(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        reset_egress_cache_for_tests()

    async def test_discovery_deduplicates_public_ip_and_keeps_bindable_source(self):
        candidates = [
            LocalAddress(None, "default", True),
            LocalAddress("10.0.0.5", "eth0"),
            LocalAddress("10.0.0.6", "eth0:1"),
        ]

        async def probe(candidate):
            public_ip = "203.0.113.10" if candidate.source_ip != "10.0.0.6" else "203.0.113.11"
            return (
                EgressChannel(
                    public_ip=public_ip,
                    source_ip=candidate.source_ip,
                    interface=candidate.interface,
                    is_default=candidate.is_default,
                ),
                "",
            )

        with (
            patch(
                "rpa_engine.egress_channels.local_address_candidates",
                return_value=candidates,
            ),
            patch(
                "rpa_engine.egress_channels._probe_local_address",
                AsyncMock(side_effect=probe),
            ),
        ):
            snapshot = await discover_egress_channels(force=True)

        self.assertEqual([item.public_ip for item in snapshot.channels], ["203.0.113.10", "203.0.113.11"])
        self.assertEqual(snapshot.channels[0].source_ip, "10.0.0.5")
        self.assertTrue(snapshot.channels[0].is_default)

    async def test_selected_channel_is_first_and_attempt_count_is_bounded(self):
        snapshot = EgressSnapshot(
            channels=(
                EgressChannel("198.51.100.1", "10.0.0.1", "eth0", True),
                EgressChannel("198.51.100.2", "10.0.0.2", "eth0:1"),
                EgressChannel("198.51.100.3", "10.0.0.3", "eth0:2"),
            ),
            errors=(),
            detected_at=time.time(),
        )
        with patch(
            "rpa_engine.egress_channels.discover_egress_channels",
            AsyncMock(return_value=snapshot),
        ):
            routes = await resolve_send_channels("198.51.100.2", 2)

        self.assertEqual([item.public_ip for item in routes], ["198.51.100.2", "198.51.100.1"])

    async def test_missing_selected_channel_fails_closed(self):
        snapshot = EgressSnapshot(
            channels=(EgressChannel("198.51.100.1", None, "default", True),),
            errors=(),
            detected_at=time.time(),
        )
        with patch(
            "rpa_engine.egress_channels.discover_egress_channels",
            AsyncMock(return_value=snapshot),
        ):
            with self.assertRaises(EgressChannelUnavailable):
                await resolve_send_channels("198.51.100.99", 2)

    async def test_browser_proxy_binds_selected_source_address(self):
        observed_peer = asyncio.get_running_loop().create_future()

        async def target_handler(reader, writer):
            if not observed_peer.done():
                observed_peer.set_result(writer.get_extra_info("peername")[0])
            payload = await reader.readexactly(4)
            writer.write(payload)
            await writer.drain()
            writer.close()
            await writer.wait_closed()

        target = await asyncio.start_server(target_handler, "127.0.0.1", 0)
        target_port = target.sockets[0].getsockname()[1]
        proxy = await SourceBoundProxy("127.0.0.2").start()
        writer = None
        try:
            reader, writer = await asyncio.open_connection(
                "127.0.0.1",
                int(proxy.server_url.rpartition(":")[2]),
            )
            writer.write(
                (
                    f"CONNECT 127.0.0.1:{target_port} HTTP/1.1\r\n"
                    f"Host: 127.0.0.1:{target_port}\r\n\r\n"
                ).encode("ascii")
            )
            await writer.drain()
            response = await reader.readuntil(b"\r\n\r\n")
            self.assertIn(b"200 Connection Established", response)

            writer.write(b"ping")
            await writer.drain()
            self.assertEqual(await reader.readexactly(4), b"ping")
            self.assertEqual(await asyncio.wait_for(observed_peer, 1), "127.0.0.2")
        finally:
            if writer is not None:
                writer.close()
                await writer.wait_closed()
            await proxy.close()
            target.close()
            await target.wait_closed()


class EgressMigrationTests(unittest.TestCase):
    def test_mysql_accounts_uses_longtext_for_browser_payloads(self):
        ddl = str(CreateTable(Account.__table__).compile(dialect=mysql.dialect()))

        self.assertIn("cookie_data LONGTEXT", ddl)
        self.assertIn("im_session_data LONGTEXT", ddl)
        self.assertIn("qr_code_base64 LONGTEXT", ddl)

    def test_old_accounts_table_receives_egress_columns(self):
        engine = create_engine("sqlite:///:memory:")
        with engine.begin() as connection:
            connection.execute(text("CREATE TABLE accounts (id INTEGER PRIMARY KEY)"))
            migrate_accounts_table(connection)
            columns = {item["name"] for item in inspect(connection).get_columns("accounts")}

        self.assertIn("egress_public_ip", columns)
        self.assertIn("egress_auto_attempts", columns)


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