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
)


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"])