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