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