clickstream-data-platform/tests/dag-probes-unit.py

"""Малые проверки логики пробника ClickHouse без запуска Airflow."""

from __future__ import annotations

import importlib.util
import sys
import types
import unittest
from pathlib import Path


class DeclaredTask:
    """Заглушка объявленной задачи: держит только цепочку через `>>`."""

    def __rshift__(self, other):
        return other


def load_clickhouse_dag_tasks():
    airflow_module = types.ModuleType("airflow")
    sdk_module = types.ModuleType("airflow.sdk")
    captured_tasks = {}

    def dag(**_kwargs):
        def decorate(function):
            return function

        return decorate

    def task(function=None, **_kwargs):
        def capture(target):
            captured_tasks[target.__name__] = target

            def declare_task(*_args, **_kwargs):
                return DeclaredTask()

            return declare_task

        return capture(function) if function is not None else capture

    sdk_module.Connection = object
    sdk_module.dag = dag
    sdk_module.get_current_context = lambda: {}
    sdk_module.task = task
    airflow_module.sdk = sdk_module
    sys.modules["airflow"] = airflow_module
    sys.modules["airflow.sdk"] = sdk_module

    dag_path = Path(__file__).resolve().parents[1] / "dags" / "test_clickhouse.py"
    spec = importlib.util.spec_from_file_location("test_clickhouse_dag", dag_path)
    if spec is None or spec.loader is None:
        raise RuntimeError("не удалось загрузить модуль пробника ClickHouse")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module, captured_tasks


def load_kafka_dag_task():
    airflow_module = types.ModuleType("airflow")
    sdk_module = types.ModuleType("airflow.sdk")
    captured_tasks = {}

    def dag(**_kwargs):
        def decorate(function):
            return function

        return decorate

    def task(function):
        captured_tasks[function.__name__] = function

        def declare_task(*_args, **_kwargs):
            return None

        return declare_task

    sdk_module.dag = dag
    sdk_module.get_current_context = lambda: {"run_id": "unit-test"}
    sdk_module.task = task
    airflow_module.sdk = sdk_module
    sys.modules["airflow"] = airflow_module
    sys.modules["airflow.sdk"] = sdk_module

    dag_path = Path(__file__).resolve().parents[1] / "dags" / "test_kafka.py"
    spec = importlib.util.spec_from_file_location("test_kafka_dag", dag_path)
    if spec is None or spec.loader is None:
        raise RuntimeError("не удалось загрузить модуль пробника Kafka")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return captured_tasks["check_round_trip"]


class QueryResult:
    def __init__(self, result_rows: list[tuple]) -> None:
        self.result_rows = result_rows


class RecordingClient:
    """Клиент ClickHouse, который запоминает запросы и отвечает заготовкой."""

    def __init__(self, remaining_tables: list[tuple[str, str]] | None = None) -> None:
        self.commands: list[str] = []
        self.queries: list[str] = []
        self.closed = False
        self._remaining_tables = remaining_tables or []

    def command(self, sql: str) -> None:
        self.commands.append(" ".join(sql.split()))

    def query(self, sql: str, parameters=None) -> QueryResult:
        self.queries.append(" ".join(sql.split()))
        return QueryResult(list(self._remaining_tables))

    def close(self) -> None:
        self.closed = True


class ClickHouseProbeTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls) -> None:
        cls.module, cls.tasks = load_clickhouse_dag_tasks()

    def run_cleanup_task(self, client: RecordingClient) -> None:
        original_client_factory = self.module._clickhouse_client
        self.module._clickhouse_client = lambda: client
        try:
            self.tasks["cleanup_tables"]()
        finally:
            self.module._clickhouse_client = original_client_factory

    def test_marker_path_requires_different_nodes_and_first_shard(self) -> None:
        marker = "свой-маркер"
        self.module._assert_marker_path(
            distributed_rows=[(1, "clickhouse-01-host", marker)],
            write_hostname="clickhouse-01-host",
            node_2_hostname="clickhouse-02-host",
            marker=marker,
        )

        with self.assertRaisesRegex(RuntimeError, "разных нод"):
            self.module._assert_marker_path(
                distributed_rows=[(2, "clickhouse-02-host", marker)],
                write_hostname="clickhouse-02-host",
                node_2_hostname="clickhouse-02-host",
                marker=marker,
            )
        with self.assertRaisesRegex(RuntimeError, "первого шарда"):
            self.module._assert_marker_path(
                distributed_rows=[(2, "clickhouse-01-host", marker)],
                write_hostname="clickhouse-01-host",
                node_2_hostname="clickhouse-02-host",
                marker=marker,
            )

    def test_cleanup_drops_tables_and_checks_both_nodes(self) -> None:
        client = RecordingClient()

        self.run_cleanup_task(client)

        dropped = {
            table
            for table in (self.module.LOCAL_TABLE, self.module.DISTRIBUTED_TABLE)
            if any(
                command.startswith(f"DROP TABLE IF EXISTS default.{table} ")
                for command in client.commands
            )
        }
        self.assertEqual(
            dropped, {self.module.LOCAL_TABLE, self.module.DISTRIBUTED_TABLE}
        )
        self.assertEqual(len(client.queries), len(self.module.NODES))
        self.assertTrue(client.closed)

    def test_cleanup_closes_client_when_tables_survive(self) -> None:
        client = RecordingClient(remaining_tables=[("airflow_probe_local", "Log")])

        with self.assertRaisesRegex(RuntimeError, "служебные таблицы остались"):
            self.run_cleanup_task(client)

        self.assertTrue(client.closed)


class KafkaProbeTests(unittest.TestCase):
    def test_producer_flushes_when_consumer_creation_fails(self) -> None:
        kafka_module = types.ModuleType("confluent_kafka")
        flush_timeouts = []

        class Message:
            def partition(self) -> int:
                return 0

            def offset(self) -> int:
                return 1

        class Producer:
            def __init__(self, _config) -> None:
                pass

            def produce(self, _topic, **kwargs) -> None:
                kwargs["on_delivery"](None, Message())

            def flush(self, timeout: int) -> int:
                flush_timeouts.append(timeout)
                return 0

        class Consumer:
            def __init__(self, _config) -> None:
                raise RuntimeError("чтение недоступно")

        kafka_module.Consumer = Consumer
        kafka_module.Producer = Producer
        kafka_module.TopicPartition = object
        sys.modules["confluent_kafka"] = kafka_module
        check_round_trip = load_kafka_dag_task()

        with self.assertRaisesRegex(RuntimeError, "чтение недоступно"):
            check_round_trip()

        self.assertEqual(flush_timeouts, [10, 1])


def run_tests() -> int:
    suite = unittest.defaultTestLoader.loadTestsFromModule(sys.modules[__name__])
    result = unittest.TestResult()
    suite.run(result)

    problems = result.failures + result.errors
    for test, details in problems:
        print(f"ОШИБКА: {test.id()}", file=sys.stderr)
        print(details, file=sys.stderr)
    passed = result.testsRun - len(problems) - len(result.skipped)
    print(f"ИТОГ: пройдено {passed}, ошибок {len(problems)}")
    return 0 if result.wasSuccessful() else 1


if __name__ == "__main__":
    raise SystemExit(run_tests())