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