140 lines
5.2 KiB
Python
140 lines
5.2 KiB
Python
"""
|
|
Unit Tests for GreatExpectationsLakehouseOperator
|
|
Verifies declarative rules (not_null, unique, in_set, row_count),
|
|
custom SQL assertions, error handling, and XCom payload formatting.
|
|
"""
|
|
import os
|
|
import sys
|
|
import pytest
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
# Ensure airflow/dags directory is on python path
|
|
DAGS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "dags"))
|
|
if DAGS_DIR not in sys.path:
|
|
sys.path.insert(0, DAGS_DIR)
|
|
|
|
from helpers.ge_lakehouse_operator import GreatExpectationsLakehouseOperator
|
|
|
|
try:
|
|
from airflow.exceptions import AirflowException
|
|
except ImportError:
|
|
class AirflowException(Exception):
|
|
pass
|
|
|
|
|
|
class TestGreatExpectationsLakehouseOperator:
|
|
|
|
@patch("helpers.ge_lakehouse_operator.create_engine")
|
|
def test_declarative_rules_all_pass(self, mock_create_engine):
|
|
mock_conn = MagicMock()
|
|
mock_engine = MagicMock()
|
|
mock_engine.connect.return_value.__enter__.return_value = mock_conn
|
|
mock_create_engine.return_value = mock_engine
|
|
|
|
# Mock query outputs:
|
|
# 1. row count = 10 (valid for min 1)
|
|
# 2. null count = 0 (valid for not_null)
|
|
# 3. in_set invalid count = 0 (valid for in_set)
|
|
mock_conn.execute.return_value.scalar.side_effect = [10, 0, 0]
|
|
|
|
op = GreatExpectationsLakehouseOperator(
|
|
task_id="test_ge_pass",
|
|
expectation_suite="permits_test_suite",
|
|
table_name="bronze_lakehouse.permits",
|
|
rules=[
|
|
{"expect": "table_row_count_to_be_between", "min_value": 1},
|
|
{"expect": "column_values_to_not_be_null", "column": "permit_id"},
|
|
{"expect": "column_values_to_be_in_set", "column": "status", "allowed_values": ["APPROVED", "REJECTED"]}
|
|
],
|
|
database_url="sqlite:///:memory:"
|
|
)
|
|
|
|
res = op.execute(context={})
|
|
assert res["status"] == "PASSED"
|
|
assert res["overall_success"] is True
|
|
assert res["total_evaluations"] == 3
|
|
assert res["successful_evaluations"] == 3
|
|
assert res["success_percent"] == 100.0
|
|
|
|
@patch("helpers.ge_lakehouse_operator.create_engine")
|
|
def test_custom_sql_query_pass(self, mock_create_engine):
|
|
mock_conn = MagicMock()
|
|
mock_engine = MagicMock()
|
|
mock_engine.connect.return_value.__enter__.return_value = mock_conn
|
|
mock_create_engine.return_value = mock_engine
|
|
|
|
# Query returns 0 (expected 0)
|
|
mock_conn.execute.return_value.scalar.return_value = 0
|
|
|
|
op = GreatExpectationsLakehouseOperator(
|
|
task_id="test_custom_sql",
|
|
expectation_suite="permits_sql_suite",
|
|
queries=[
|
|
{
|
|
"name": "check_no_negative_amounts",
|
|
"sql": "SELECT count(*) FROM table WHERE amount < 0",
|
|
"expected": 0,
|
|
"error_message": "Found negative amounts!"
|
|
}
|
|
],
|
|
database_url="sqlite:///:memory:"
|
|
)
|
|
|
|
res = op.execute(context={})
|
|
assert res["status"] == "PASSED"
|
|
assert res["overall_success"] is True
|
|
assert len(res["details"]) == 1
|
|
assert res["details"][0]["observed_value"] == 0
|
|
|
|
@patch("helpers.ge_lakehouse_operator.create_engine")
|
|
def test_fail_fast_raises_airflow_exception(self, mock_create_engine):
|
|
mock_conn = MagicMock()
|
|
mock_engine = MagicMock()
|
|
mock_engine.connect.return_value.__enter__.return_value = mock_conn
|
|
mock_create_engine.return_value = mock_engine
|
|
|
|
# Mock: 3 null values found in permit_id
|
|
mock_conn.execute.return_value.scalar.return_value = 3
|
|
|
|
op = GreatExpectationsLakehouseOperator(
|
|
task_id="test_fail_fast",
|
|
expectation_suite="permits_fail_suite",
|
|
table_name="bronze_lakehouse.permits",
|
|
rules=[
|
|
{"expect": "column_values_to_not_be_null", "column": "permit_id"}
|
|
],
|
|
fail_on_error=True,
|
|
database_url="sqlite:///:memory:"
|
|
)
|
|
|
|
with pytest.raises(Exception) as excinfo:
|
|
op.execute(context={})
|
|
assert "Great Expectations Quality Gate GAGAL" in str(excinfo.value)
|
|
|
|
@patch("helpers.ge_lakehouse_operator.create_engine")
|
|
def test_fail_on_error_false_does_not_raise(self, mock_create_engine):
|
|
mock_conn = MagicMock()
|
|
mock_engine = MagicMock()
|
|
mock_engine.connect.return_value.__enter__.return_value = mock_conn
|
|
mock_create_engine.return_value = mock_engine
|
|
|
|
# Mock: 2 invalid records
|
|
mock_conn.execute.return_value.scalar.return_value = 2
|
|
|
|
op = GreatExpectationsLakehouseOperator(
|
|
task_id="test_warn_only",
|
|
expectation_suite="permits_warn_suite",
|
|
table_name="silver.permits",
|
|
rules=[
|
|
{"expect": "column_values_to_be_in_set", "column": "status", "allowed_values": ["APPROVED"]}
|
|
],
|
|
fail_on_error=False,
|
|
database_url="sqlite:///:memory:"
|
|
)
|
|
|
|
res = op.execute(context={})
|
|
assert res["status"] == "FAILED"
|
|
assert res["overall_success"] is False
|
|
assert res["unsuccessful_evaluations"] == 1
|
|
assert res["total_unexpected_records"] == 2
|