204 lines
7.6 KiB
Python
204 lines
7.6 KiB
Python
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"])
|