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

"""Малые проверки логики пробников ClickHouse и Kafka без запуска 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_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):
        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


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)


DELIVERED_PARTITION = 3
DELIVERED_OFFSET = 42


class StubMessage:
    """Сообщение Kafka в том объёме, в каком его читает пробник."""

    def __init__(self, value: bytes) -> None:
        self._value = value

    def partition(self) -> int:
        return DELIVERED_PARTITION

    def offset(self) -> int:
        return DELIVERED_OFFSET

    def error(self):
        return None

    def value(self) -> bytes:
        return self._value


class StubTopicPartition:
    def __init__(self, topic: str, partition: int, offset: int) -> None:
        self.topic = topic
        self.partition = partition
        self.offset = offset


def install_kafka_stub(produce_error: str | None = None):
    """Ставит заглушку `confluent_kafka` и возвращает журналы её клиентов.

    Продюсер подтверждает доставку сразу и по известному адресу, консьюмер
    отдаёт записанное с первого опроса. Если задан `produce_error`, запись
    падает — так проверяется, что продюсер закрывается и на пути отказа.
    """
    producers = []
    consumers = []

    class Producer:
        def __init__(self, _config) -> None:
            self.written = b""
            self.flushed = False
            producers.append(self)

        def produce(self, _topic, key=None, value=None, on_delivery=None) -> None:
            if produce_error is not None:
                raise RuntimeError(produce_error)
            self.written = value
            on_delivery(None, StubMessage(value))

        def flush(self, _timeout) -> int:
            self.flushed = True
            return 0

    class Consumer:
        def __init__(self, _config) -> None:
            self.assigned = []
            self.closed = False
            consumers.append(self)

        def assign(self, partitions) -> None:
            self.assigned = partitions

        def poll(self, _timeout):
            return StubMessage(producers[-1].written)

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

    kafka_module = types.ModuleType("confluent_kafka")
    kafka_module.Consumer = Consumer
    kafka_module.Producer = Producer
    kafka_module.TopicPartition = StubTopicPartition
    sys.modules["confluent_kafka"] = kafka_module
    return producers, consumers


class KafkaProbeTests(unittest.TestCase):
    def test_producer_is_closed_when_write_fails(self) -> None:
        producers, _ = install_kafka_stub(produce_error="брокер недоступен")
        tasks = load_kafka_dag_tasks()

        with self.assertRaisesRegex(RuntimeError, "брокер недоступен"):
            tasks["write_marker"]()

        self.assertEqual(len(producers), 1)
        self.assertTrue(producers[0].flushed)

    def test_read_goes_to_the_address_broker_returned(self) -> None:
        _, consumers = install_kafka_stub()
        tasks = load_kafka_dag_tasks()

        written = tasks["write_marker"]()
        tasks["read_marker"](written)

        self.assertEqual(
            (written["partition"], written["offset"]),
            (DELIVERED_PARTITION, DELIVERED_OFFSET),
        )
        self.assertEqual(len(consumers), 1)
        self.assertEqual(len(consumers[0].assigned), 1)
        assigned = consumers[0].assigned[0]
        self.assertEqual(assigned.partition, DELIVERED_PARTITION)
        self.assertEqual(assigned.offset, DELIVERED_OFFSET)
        self.assertTrue(consumers[0].closed)


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