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


def load_clickhouse_dag_module():
    airflow_module = types.ModuleType("airflow")
    sdk_module = types.ModuleType("airflow.sdk")

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

        return decorate

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

        return declare_task

    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


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 FailingCleanupClient:
    def __init__(self) -> None:
        self.closed = False

    def command(self, _sql: str) -> None:
        raise RuntimeError("очистка недоступна")

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


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

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

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

    def test_cleanup_keeps_original_error_as_primary(self) -> None:
        client = FailingCleanupClient()
        original_error = RuntimeError("маркер не найден")

        self.module._cleanup_clickhouse_client(client, original_error)

        self.assertTrue(client.closed)
        self.assertIn("очистка недоступна", "\n".join(original_error.__notes__))


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