airflow-manual/airflow-docker/tests/test_greenplum_helpers.py

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)