Первая публикация кода

2025-10-17 19:16:15 +03:00
commit 0dffd5df20
17 changed files with 3440 additions and 0 deletions
@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import importlib
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import Type
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.append(str(PROJECT_ROOT))
+
+if "airflow" not in sys.modules:
+    airflow_module = ModuleType("airflow")
+    airflow_module.__path__ = [str(PROJECT_ROOT / "airflow")]
+    sys.modules["airflow"] = airflow_module
+
+    providers_module = ModuleType("airflow.providers")
+    providers_module.__path__ = []
+    sys.modules["airflow.providers"] = providers_module
+    airflow_module.providers = providers_module
+
+    postgres_module = ModuleType("airflow.providers.postgres")
+    postgres_module.__path__ = []
+    sys.modules["airflow.providers.postgres"] = postgres_module
+    providers_module.postgres = postgres_module
+
+    hooks_module = ModuleType("airflow.providers.postgres.hooks")
+    hooks_module.__path__ = []
+    sys.modules["airflow.providers.postgres.hooks"] = hooks_module
+    postgres_module.hooks = hooks_module
+
+if "psycopg2" not in sys.modules:
+    psycopg2_stub = ModuleType("psycopg2")
+    psycopg2_stub.connect = lambda **_: None  # type: ignore[assignment]
+    sys.modules["psycopg2"] = psycopg2_stub
+
+
+def _ensure_stub_module(full_name: str) -> ModuleType:
+    """
+    Ensure that module placeholders exist for a dotted path and return leaf module.
+    """
+    parts = full_name.split(".")
+    module: ModuleType | None = None
+    path = ""
+    for part in parts:
+        path = f"{path}.{part}" if path else part
+        if path not in sys.modules:
+            new_module = ModuleType(path)
+            if module is not None:
+                setattr(module, part, new_module)
+            sys.modules[path] = new_module
+            module = new_module
+        else:
+            module = sys.modules[path]
+    assert isinstance(module, ModuleType)
+    return module
+
+
+def patch_postgres_hook(monkeypatch, hook_cls: Type) -> None:
+    """
+    Patch PostgresHook so that helpers.greenplum can be exercised without real Airflow.
+    """
+    try:
+        module = importlib.import_module("airflow.providers.postgres.hooks.postgres")
+    except ModuleNotFoundError:
+        module = _ensure_stub_module("airflow.providers.postgres.hooks.postgres")
+    monkeypatch.setattr(module, "PostgresHook", hook_cls, raising=False)
@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import importlib
+
+import pytest
+
+
+def _airflow_available() -> bool:
+    try:
+        af = importlib.import_module("airflow")
+    except Exception:
+        return False
+    # Real Airflow exposes DAG at top-level
+    return hasattr(af, "DAG")
+
+
+pytestmark = pytest.mark.skipif(not _airflow_available(), reason="Airflow is not installed for DAG smoke tests")
+
+
+def _load_dag(module_name: str):
+    mod = importlib.import_module(module_name)
+    assert hasattr(mod, "dag"), f"{module_name} must expose variable 'dag'"
+    return getattr(mod, "dag")
+
+
+def test_csv_to_greenplum_dag_structure():
+    dag = _load_dag("airflow.dags.csv_to_greenplum")
+
+    # tasks
+    expected_tasks = {
+        "create_orders_table",
+        "generate_csv",
+        "preview_csv",
+        "load_csv_to_greenplum",
+    }
+    assert expected_tasks.issubset(dag.task_dict.keys())
+
+    # linear dependencies
+    t1 = dag.get_task("create_orders_table")
+    t2 = dag.get_task("generate_csv")
+    t3 = dag.get_task("preview_csv")
+    t4 = dag.get_task("load_csv_to_greenplum")
+
+    assert t2 in t1.get_direct_relatives("downstream")
+    assert t3 in t2.get_direct_relatives("downstream")
+    assert t4 in t3.get_direct_relatives("downstream")
+
+
+def test_data_quality_greenplum_dag_structure():
+    dag = _load_dag("airflow.dags.data_quality_greenplum")
+
+    expected_tasks = {
+        "check_orders_table_exists",
+        "check_orders_schema",
+        "check_orders_has_rows",
+        "check_order_duplicates",
+        "data_quality_summary",
+    }
+    assert expected_tasks.issubset(dag.task_dict.keys())
+
+    e = dag.get_task("check_orders_table_exists")
+    s = dag.get_task("check_orders_schema")
+    h = dag.get_task("check_orders_has_rows")
+    d = dag.get_task("check_order_duplicates")
+    q = dag.get_task("data_quality_summary")
+
+    assert s in e.get_direct_relatives("downstream")
+    assert h in s.get_direct_relatives("downstream")
+    assert d in h.get_direct_relatives("downstream")
+    assert q in d.get_direct_relatives("downstream")
@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, List, Sequence
+
+import pytest
+
+import airflow.dags.helpers.greenplum as greenplum
+from tests.conftest import patch_postgres_hook
+
+
+@dataclass
+class FakeCursor:
+    fetchone_value: Any = None
+    fetchall_value: Sequence[Any] | None = None
+    rowcount: int | None = None
+
+    def __post_init__(self) -> None:
+        self.queries: List[Any] = []
+
+    def execute(self, query: str, params: Any | None = None) -> None:
+        self.queries.append((query, params))
+
+    def fetchone(self) -> Any:
+        return self.fetchone_value
+
+    def fetchall(self) -> Sequence[Any] | None:
+        return self.fetchall_value
+
+    def __enter__(self) -> FakeCursor:
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        return None
+
+
+class FakeConn:
+    def __init__(self, cursors: Sequence[FakeCursor]) -> None:
+        self._cursors = list(cursors)
+        self._index = 0
+        self.commits = 0
+
+    def cursor(self) -> FakeCursor:
+        cursor = self._cursors[self._index]
+        self._index += 1
+        return cursor
+
+    def commit(self) -> None:
+        self.commits += 1
+
+
+def test_get_gp_conn_uses_airflow_hook(monkeypatch) -> None:
+    class FakeHook:
+        def __init__(self, postgres_conn_id: str) -> None:
+            self.postgres_conn_id = postgres_conn_id
+
+        def get_conn(self) -> str:
+            return "hook_connection"
+
+    patch_postgres_hook(monkeypatch, FakeHook)
+    monkeypatch.setattr(greenplum, "GP_CONN_ID", "demo_conn", raising=False)
+    monkeypatch.setattr(greenplum, "GP_USE_AIRFLOW_CONN", True, raising=False)
+
+    conn = greenplum.get_gp_conn()
+
+    assert conn == "hook_connection"
+
+
+def test_get_gp_conn_fallback_to_psycopg(monkeypatch) -> None:
+    class BrokenHook:
+        def __init__(self, postgres_conn_id: str) -> None:
+            self.postgres_conn_id = postgres_conn_id
+
+        def get_conn(self):
+            raise RuntimeError("boom")
+
+    patch_postgres_hook(monkeypatch, BrokenHook)
+    monkeypatch.setattr(greenplum, "GP_USE_AIRFLOW_CONN", True, raising=False)
+    monkeypatch.setattr(greenplum, "GP_CONN_ID", "demo_conn", raising=False)
+    monkeypatch.setenv("GP_DB", "demo_db")
+    monkeypatch.setenv("GP_USER", "demo_user")
+    monkeypatch.setenv("GP_PASSWORD", "secret")
+    monkeypatch.setenv("GP_HOST", "greenplum-host")
+    monkeypatch.setenv("GP_PORT", "5434")
+
+    captured_kwargs = {}
+
+    def fake_connect(**kwargs):
+        captured_kwargs.update(kwargs)
+        return "psycopg_connection"
+
+    monkeypatch.setattr(greenplum.psycopg2, "connect", fake_connect)
+
+    conn = greenplum.get_gp_conn()
+
+    assert conn == "psycopg_connection"
+    assert captured_kwargs == {
+        "dbname": "demo_db",
+        "user": "demo_user",
+        "password": "secret",
+        "host": "greenplum-host",
+        "port": 5434,
+    }
+
+
+def test_get_gp_conn_without_airflow(monkeypatch) -> None:
+    monkeypatch.setattr(greenplum, "GP_USE_AIRFLOW_CONN", False, raising=False)
+    monkeypatch.setenv("GP_DB", "demo_db")
+    monkeypatch.setenv("GP_USER", "demo_user")
+    monkeypatch.setenv("GP_PASSWORD", "secret")
+    monkeypatch.setenv("GP_HOST", "greenplum-host")
+    monkeypatch.setenv("GP_PORT", "5435")
+
+    captured_kwargs = {}
+
+    def fake_connect(**kwargs):
+        captured_kwargs.update(kwargs)
+        return "direct_psycopg"
+
+    monkeypatch.setattr(greenplum.psycopg2, "connect", fake_connect)
+
+    conn = greenplum.get_gp_conn()
+
+    assert conn == "direct_psycopg"
+    assert captured_kwargs["port"] == 5435
+
+
+def test_assert_orders_table_exists_ok() -> None:
+    conn = FakeConn([FakeCursor(fetchone_value=(1,))])
+
+    greenplum.assert_orders_table_exists(conn)
+
+
+def test_assert_orders_table_exists_missing() -> None:
+    conn = FakeConn([FakeCursor(fetchone_value=None)])
+
+    with pytest.raises(ValueError):
+        greenplum.assert_orders_table_exists(conn)
+
+
+def test_assert_orders_schema_ok() -> None:
+    expected = list(greenplum.EXPECTED_ORDERS_SCHEMA)
+    conn = FakeConn([FakeCursor(fetchall_value=expected)])
+
+    greenplum.assert_orders_schema(conn)
+
+
+def test_assert_orders_schema_mismatch() -> None:
+    conn = FakeConn([FakeCursor(fetchall_value=[("order_id", "bigint")])])
+
+    with pytest.raises(ValueError):
+        greenplum.assert_orders_schema(conn)
+
+
+def test_assert_orders_have_rows_ok() -> None:
+    conn = FakeConn([FakeCursor(fetchone_value=(5,))])
+
+    greenplum.assert_orders_have_rows(conn)
+
+
+def test_assert_orders_have_rows_empty() -> None:
+    conn = FakeConn([FakeCursor(fetchone_value=(0,))])
+
+    with pytest.raises(ValueError):
+        greenplum.assert_orders_have_rows(conn)
+
+
+def test_assert_orders_no_duplicates_ok() -> None:
+    conn = FakeConn([FakeCursor(fetchone_value=(0,))])
+
+    greenplum.assert_orders_no_duplicates(conn)
+
+
+def test_assert_orders_no_duplicates_detected() -> None:
+    conn = FakeConn([FakeCursor(fetchone_value=(3,))])
+
+    with pytest.raises(ValueError):
+        greenplum.assert_orders_no_duplicates(conn)