89 lines
3.8 KiB
Python
89 lines
3.8 KiB
Python
|
import os
|
||
|
|
import sys
|
||
|
|
import unittest
|
||
|
|
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.base_lakehouse_operator import BaseLakehouseOperator
|
||
|
|
from helpers.duckdb_operator import DuckDBIcebergOperator
|
||
|
|
from helpers.pyiceberg_operator import PyIcebergTransformationOperator
|
||
|
|
from helpers.dbt_trino_operator import DbtTrinoIcebergOperator
|
||
|
|
from helpers.pyspark_operator import PySparkIcebergOperator
|
||
|
|
|
||
|
|
|
||
|
|
class TestLakehouseOperators(unittest.TestCase):
|
||
|
|
def test_base_lakehouse_operator_initialization(self):
|
||
|
|
op = BaseLakehouseOperator(
|
||
|
|
task_id="test_base_task",
|
||
|
|
inlets=["s3://dlh-bronze/test_data.parquet"],
|
||
|
|
outlets=["iceberg://dlh_lakehouse.silver.test_data"],
|
||
|
|
pipeline_id="test_pipeline",
|
||
|
|
layer="silver",
|
||
|
|
engine="generic"
|
||
|
|
)
|
||
|
|
self.assertEqual(op.task_id, "test_base_task")
|
||
|
|
self.assertEqual(op.inlets, ["s3://dlh-bronze/test_data.parquet"])
|
||
|
|
self.assertEqual(op.outlets, ["iceberg://dlh_lakehouse.silver.test_data"])
|
||
|
|
self.assertEqual(op.layer, "silver")
|
||
|
|
self.assertEqual(op.engine, "generic")
|
||
|
|
|
||
|
|
def test_duckdb_iceberg_operator_initialization(self):
|
||
|
|
op = DuckDBIcebergOperator(
|
||
|
|
task_id="silver_duckdb_task",
|
||
|
|
source_parquet_path="s3://dlh-bronze/environmental_permits.parquet",
|
||
|
|
target_table_name="dlh_lakehouse.silver.environmental_permits",
|
||
|
|
primary_keys=["id", "updated_at"]
|
||
|
|
)
|
||
|
|
self.assertEqual(op.task_id, "silver_duckdb_task")
|
||
|
|
self.assertEqual(op.inlets, ["s3://dlh-bronze/environmental_permits.parquet"])
|
||
|
|
self.assertEqual(op.outlets, ["dlh_lakehouse.silver.environmental_permits"])
|
||
|
|
self.assertEqual(op.layer, "silver")
|
||
|
|
self.assertEqual(op.engine, "duckdb")
|
||
|
|
self.assertIn("PARTITION BY id", op.transformation_sql)
|
||
|
|
|
||
|
|
def test_pyiceberg_operator_initialization(self):
|
||
|
|
op = PyIcebergTransformationOperator(
|
||
|
|
task_id="silver_pyiceberg_task",
|
||
|
|
source_parquet_path="s3://dlh-bronze/waste_reports.parquet",
|
||
|
|
target_table_name="dlh_lakehouse.silver.waste_reports",
|
||
|
|
write_mode="overwrite"
|
||
|
|
)
|
||
|
|
self.assertEqual(op.task_id, "silver_pyiceberg_task")
|
||
|
|
self.assertEqual(op.inlets, ["s3://dlh-bronze/waste_reports.parquet"])
|
||
|
|
self.assertEqual(op.outlets, ["dlh_lakehouse.silver.waste_reports"])
|
||
|
|
self.assertEqual(op.layer, "silver")
|
||
|
|
self.assertEqual(op.engine, "pyiceberg")
|
||
|
|
self.assertEqual(op.write_mode, "overwrite")
|
||
|
|
|
||
|
|
def test_dbt_trino_operator_initialization(self):
|
||
|
|
op = DbtTrinoIcebergOperator(
|
||
|
|
task_id="gold_dbt_task",
|
||
|
|
select_models="gold_permits_summary_daily",
|
||
|
|
layer="gold"
|
||
|
|
)
|
||
|
|
self.assertEqual(op.task_id, "gold_dbt_task")
|
||
|
|
self.assertEqual(op.select_models, "gold_permits_summary_daily")
|
||
|
|
self.assertEqual(op.layer, "gold")
|
||
|
|
self.assertEqual(op.engine, "dbt_trino")
|
||
|
|
self.assertEqual(op.outlets, ["gold.gold_permits_summary_daily"])
|
||
|
|
|
||
|
|
def test_pyspark_operator_initialization(self):
|
||
|
|
op = PySparkIcebergOperator(
|
||
|
|
task_id="silver_spark_task",
|
||
|
|
source_parquet_path="s3://dlh-bronze/spku_realtime.parquet",
|
||
|
|
target_table_name="dlh_lakehouse.silver.spku_realtime"
|
||
|
|
)
|
||
|
|
self.assertEqual(op.task_id, "silver_spark_task")
|
||
|
|
self.assertEqual(op.inlets, ["s3://dlh-bronze/spku_realtime.parquet"])
|
||
|
|
self.assertEqual(op.outlets, ["dlh_lakehouse.silver.spku_realtime"])
|
||
|
|
self.assertEqual(op.layer, "silver")
|
||
|
|
self.assertEqual(op.engine, "pyspark")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|