merica-pipelines/tests/test_helpers.py

import os
import sys
import json
import pytest
import pandas as pd
from datetime import datetime, timedelta

# 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.storage_helper import storage_manager, build_partitioned_path, StorageManager
from helpers.api_helper import (
    GenericAPIExtractorOperator,
    create_resilient_session,
    generate_synthetic_telemetry
)
from helpers.sling_helper import (
    SlingDBIngestionOperator,
    get_checkpoint_value,
    set_checkpoint_value,
    generate_synthetic_permits_data
)
from helpers.gws_helper import (
    GWSExtractorOperator,
    DataCleansingEngine,
    generate_synthetic_gws_waste_data
)
from helpers.dlt_helper import (
    DLTAPIExtractorOperator,
    flatten_json_record,
    normalize_api_response,
    generate_synthetic_dlt_telemetry
)


class TestStorageHelper:
    def test_build_partitioned_path(self):
        fixed_dt = datetime(2026, 9, 28, 10, 0, 0)
        path = build_partitioned_path("api/spku_api", "telemetry.json", dt=fixed_dt)
        assert path == "api/spku_api/year=2026/month=09/day=28/telemetry.json"

    def test_save_json_local_fallback(self):
        test_payload = {"service": "merica", "metric": "air_quality", "value": 42}
        result = storage_manager.save_json(
            data=test_payload,
            bucket="dlh-bronze-test",
            key="test/sample.json"
        )
        assert result["s3_uri"] == "s3://dlh-bronze-test/test/sample.json"
        assert os.path.exists(result["local_path"])
        with open(result["local_path"], "r") as f:
            loaded = json.load(f)
        assert loaded["value"] == 42
        os.remove(result["local_path"])

    def test_save_dataframe_csv(self):
        df = pd.DataFrame({"col1": [1, 2], "col2": ["A", "B"]})
        result = storage_manager.save_dataframe(
            df=df,
            bucket="dlh-bronze-test",
            key="test/table.csv",
            target_format="csv"
        )
        assert result["format"] == "csv"
        assert result["rows_written"] == 2
        assert os.path.exists(result["local_path"])
        read_df = pd.read_csv(result["local_path"])
        assert len(read_df) == 2
        os.remove(result["local_path"])


class TestApiHelper:
    def test_resilient_session(self):
        session = create_resilient_session(max_retries=2, backoff_factor=1.0)
        assert session is not None
        assert "http://" in session.adapters
        assert "https://" in session.adapters

    def test_generate_synthetic_telemetry(self):
        telemetry = generate_synthetic_telemetry("SPKU_01", "2026-09-28T00:00:00Z", "2026-09-28T03:00:00Z")
        assert telemetry["status"] == "success"
        assert telemetry["station_id"] == "SPKU_01"
        assert len(telemetry["readings"]) > 0
        assert "pm25" in telemetry["readings"][0]

    def test_api_extractor_operator_execution(self):
        config = {
            "source_type": "api",
            "connection_id": "conn_spku_api_service",
            "api_endpoint": "https://api.spku.dlh.id/v1/stasiun/spku/data-realtime",
            "lookback_hours": 3,
            "target": {
                "bucket": "dlh-bronze-test",
                "prefix": "test_api/spku"
            }
        }
        operator = GenericAPIExtractorOperator(
            task_id="test_api_op",
            bronze_config=config,
            mock_on_network_error=True
        )
        result = operator.execute(context={"execution_date": datetime(2026, 9, 28, 12, 0, 0)})
        
        assert result["status"] == "success"
        assert result["record_count"] >= 1
        assert "window_start" in result
        assert "window_end" in result
        assert os.path.exists(result["local_path"])
        os.remove(result["local_path"])


class TestSlingHelper:
    def test_checkpoint_state(self):
        key = "unit_test_checkpoint"
        set_checkpoint_value(key, "2026-09-28 10:00:00")
        val = get_checkpoint_value(key)
        assert val == "2026-09-28 10:00:00"
        state_file = f"/tmp/{key}.state"
        if os.path.exists(state_file):
            os.remove(state_file)

    def test_generate_synthetic_permits(self):
        df = generate_synthetic_permits_data(row_count=4)
        assert len(df) == 4
        assert "permit_number" in df.columns
        assert "status" in df.columns

    def test_sling_db_ingestion_operator_execution(self):
        config = {
            "source_type": "database",
            "connection_id": "conn_postgres_source_db",
            "table_name": "public.environmental_permits",
            "ingestion_mode": "incremental_watermark",
            "watermark_column": "watermark_ts",
            "target": {
                "bucket": "dlh-bronze-test",
                "format": "csv",
                "prefix": "test_db/permits"
            }
        }
        operator = SlingDBIngestionOperator(
            task_id="test_sling_op",
            bronze_config=config
        )
        result = operator.execute(context={"execution_date": datetime(2026, 9, 28, 12, 0, 0)})
        
        assert result["status"] == "success"
        assert result["rows_extracted"] > 0
        assert result["checkpoint_new"] is not None
        assert os.path.exists(result["local_path"])
        os.remove(result["local_path"])


class TestGwsHelper:
    def test_data_cleansing_headers(self):
        assert DataCleansingEngine.clean_header("Total Volume (Kg) ") == "total_volume_kg"
        assert DataCleansingEngine.clean_header("Nama TPS / Lokasi") == "nama_tps_lokasi"
        assert DataCleansingEngine.clean_header("Biaya Retribusi (Rp)") == "biaya_retribusi_rp"

    def test_data_cleansing_numeric_units(self):
        assert DataCleansingEngine.clean_numeric_value("1,500 kg") == 1500
        assert DataCleansingEngine.clean_numeric_value("Rp 250,000") == 250000
        assert DataCleansingEngine.clean_numeric_value("25.5 Ton") == 25.5
        assert DataCleansingEngine.clean_numeric_value("Jakarta Timur") == "Jakarta Timur"

    def test_data_cleansing_dataframe_sanitization(self):
        raw_df = generate_synthetic_gws_waste_data()
        cleaned_df = DataCleansingEngine.sanitize_dataframe(raw_df)
        
        # Headers sanitized
        assert "wilayah_kecamatan" in cleaned_df.columns
        assert "volume_sampah_kg" in cleaned_df.columns
        assert "biaya_retribusi_rp" in cleaned_df.columns
        
        # Merged cells forward-filled (first column had empty strings)
        assert cleaned_df["wilayah_kecamatan"].isna().sum() == 0
        assert cleaned_df["wilayah_kecamatan"].iloc[1] == "Jakarta Timur"
        assert cleaned_df["wilayah_kecamatan"].iloc[4] == "Jakarta Selatan"
        
        # Numeric values cleaned
        assert cleaned_df["volume_sampah_kg"].iloc[0] == 1250
        assert cleaned_df["biaya_retribusi_rp"].iloc[0] == 250000

    def test_gws_extractor_operator_execution(self):
        config = {
            "source_type": "gws",
            "connection_id": "conn_gws_service_account",
            "spreadsheet_id": "1OSRugXZvfjoY9M0yySKT-F3L4LTZXa40J0knU3_4Hik",
            "sheet_range": "Sheet1!A1:Z1000",
            "target": {
                "bucket": "dlh-bronze-test",
                "format": "csv",
                "prefix": "test_gws/waste_reports"
            }
        }
        operator = GWSExtractorOperator(
            task_id="test_gws_op",
            bronze_config=config,
            auto_clean=True
        )
        result = operator.execute(context={"execution_date": datetime(2026, 9, 28, 12, 0, 0)})
        
        assert result["status"] == "success"
        assert result["rows_extracted"] == 5
        assert "volume_sampah_kg" in result["columns"]
        assert os.path.exists(result["local_path"])
        os.remove(result["local_path"])


class TestDltHelper:
    def test_flatten_json_record(self):
        nested = {
            "station_id": "SPKU_01",
            "location": {"city": "Jakarta", "coords": {"lat": -6.2, "lon": 106.8}},
            "readings": {"pm25": 42.5, "pm10": 60.1},
            "tags": ["ambient", "urban"]
        }
        flat = flatten_json_record(nested)
        assert flat["station_id"] == "SPKU_01"
        assert flat["location_city"] == "Jakarta"
        assert flat["location_coords_lat"] == -6.2
        assert flat["location_coords_lon"] == 106.8
        assert flat["readings_pm25"] == 42.5
        assert flat["readings_pm10"] == 60.1
        assert "urban" in flat["tags"]

    def test_normalize_api_response_dict_and_list(self):
        # Case A: Dict with "data" key
        payload = {
            "status": "success",
            "data": [
                {"id": 1, "metrics": {"value": 100}},
                {"id": 2, "metrics": {"value": 200}}
            ]
        }
        records = normalize_api_response(payload, unnest_mode="auto_flatten")
        assert len(records) == 2
        assert records[0]["metrics_value"] == 100
        assert records[1]["metrics_value"] == 200

        # Case B: Top-level list
        list_payload = [{"device": "A", "val": 10}, {"device": "B", "val": 20}]
        rec_list = normalize_api_response(list_payload, unnest_mode="auto_flatten")
        assert len(rec_list) == 2
        assert rec_list[0]["device"] == "A"

    def test_generate_synthetic_dlt_telemetry(self):
        synth = generate_synthetic_dlt_telemetry("https://api.spku.dlh.id/v1/telemetry", limit=3)
        assert synth["status"] == "success"
        assert len(synth["data"]) == 3
        assert "station_id" in synth["data"][0]
        assert "telemetry" in synth["data"][0]

    def test_dlt_api_extractor_operator_execution(self):
        config = {
            "endpoint_url": "https://api.spku.dlh.id/v1/telemetry/hourly",
            "http_method": "GET",
            "pagination": {
                "strategy": "page_number",
                "page_param": "page",
                "limit_param": "limit",
                "page_size": 4,
                "max_pages": 1
            },
            "incremental": {
                "cursor_column": "measured_at",
                "lookback_window_hours": 3
            },
            "unnest_mode": "auto_flatten",
            "landing_zone": {
                "bucket": "dlh-landing-test",
                "prefix": "test_landing/api"
            },
            "bronze_zone": {
                "bucket": "dlh-bronze-test",
                "prefix": "test_bronze/api",
                "format": "parquet"
            }
        }
        operator = DLTAPIExtractorOperator(
            task_id="test_dlt_api_op",
            bronze_config=config,
            mock_on_network_error=True
        )
        result = operator.execute(context={"execution_date": datetime(2026, 9, 28, 12, 0, 0)})

        assert result["status"] == "success"
        assert result["rows_extracted"] >= 1
        assert "landing_uri" in result
        assert "bronze_uri" in result
        assert "telemetry_pm25" in result["columns"] or "pm25" in str(result["columns"])
        assert "_dlt_load_id" in result["columns"]
        assert "_ingested_at" in result["columns"]