Northern-Thailand-Ping-River-Monitor/rust/collector/src/db_tests.rs

//! Database tests against a real PostgreSQL. Each test runs in its own
//! throwaway schema. Set TEST_DATABASE_URL to run them, e.g.
//! `postgres://postgres:ping@localhost:55432/ping`; without it they are
//! skipped (reported as passed with a note on stderr).

use crate::dams::DamRecord;
use crate::db;
use crate::hii::{RainRecord, WaterLevelRecord};
use crate::rain::RainRow;
use crate::rid::{self, Reading, RidSource};
use crate::stations::Stations;
use anyhow::Result;
use chrono::{Duration, NaiveDate, NaiveDateTime, NaiveTime};
use sqlx::PgPool;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::Mutex;

struct TestDb {
    pool: PgPool,
    admin: PgPool,
    schema: String,
}

impl TestDb {
    async fn new() -> Option<Self> {
        let Ok(url) = std::env::var("TEST_DATABASE_URL") else {
            eprintln!("TEST_DATABASE_URL not set: skipping database test");
            return None;
        };
        let schema = format!(
            "t_{}_{}",
            std::process::id(),
            chrono::Utc::now().timestamp_nanos_opt().unwrap()
        );
        let admin = PgPool::connect(&url).await.expect("connect");
        sqlx::query(&format!("CREATE SCHEMA {schema}"))
            .execute(&admin)
            .await
            .unwrap();
        let opts = PgConnectOptions::from_str(&url)
            .unwrap()
            .options([("search_path", schema.as_str())]);
        let pool = PgPoolOptions::new()
            .max_connections(2)
            .connect_with(opts)
            .await
            .unwrap();
        db::ensure_schema(&pool).await.expect("schema");
        Some(TestDb {
            pool,
            admin,
            schema,
        })
    }

    async fn drop(self) {
        self.pool.close().await;
        sqlx::query(&format!("DROP SCHEMA {} CASCADE", self.schema))
            .execute(&self.admin)
            .await
            .unwrap();
    }
}

fn d(y: i32, m: u32, day: u32) -> NaiveDate {
    NaiveDate::from_ymd_opt(y, m, day).unwrap()
}
fn t(y: i32, m: u32, day: u32, h: u32, min: u32) -> NaiveDateTime {
    d(y, m, day).and_hms_opt(h, min, 0).unwrap()
}

/// Canned RID: date -> last published hour (1..=24); P.1 (column 8) and P.20
/// (column 1). Records which reports were requested.
struct FakeRid {
    published: HashMap<NaiveDate, i64>,
    calls: Mutex<Vec<NaiveDate>>,
}

impl FakeRid {
    fn new(published: &[(NaiveDate, i64)]) -> Self {
        FakeRid {
            published: published.iter().copied().collect(),
            calls: Mutex::new(Vec::new()),
        }
    }
    fn calls(&self) -> Vec<NaiveDate> {
        self.calls.lock().unwrap().clone()
    }
}

impl RidSource for FakeRid {
    async fn fetch_report(&self, date: NaiveDate) -> Result<Vec<Reading>> {
        self.calls.lock().unwrap().push(date);
        let last = self.published.get(&date).copied().unwrap_or(0);
        Ok((1..=last)
            .flat_map(|h| {
                let ts = date.and_time(NaiveTime::MIN) + Duration::hours(h);
                [8u32, 1].map(|n| Reading {
                    timestamp: ts,
                    station_num: n,
                    water_level: 1.0 + h as f64 / 100.0,
                    discharge: Some(50.0),
                    discharge_percent: Some(12.5),
                })
            })
            .collect())
    }
}

async fn stored_hours(pool: &PgPool, day: NaiveDate) -> Vec<u32> {
    let rec = db::recorded_hours_by_day(pool, day, day).await.unwrap();
    let mut h: Vec<u32> = rec
        .get(&day)
        .cloned()
        .unwrap_or_default()
        .into_iter()
        .collect();
    h.sort();
    h
}

#[tokio::test]
async fn late_report_is_stored_and_reported_behind() {
    let Some(db) = TestDb::new().await else {
        return;
    };
    let today = d(2026, 10, 17);
    let yesterday = today - Duration::days(1);
    let stations = Stations::bundled();
    // yesterday complete already
    let full = FakeRid::new(&[(yesterday, 24)]);
    rid::fill_gaps(&full, &db.pool, &stations, Some(1), t(2026, 10, 17, 0, 30))
        .await
        .unwrap();

    // 21:30 and RID has published up to 20:00
    let src = FakeRid::new(&[(today, 20)]);
    let o = rid::run_cycle(&src, &db.pool, &stations, t(2026, 10, 17, 21, 30)).await;
    assert!(!o.current);
    assert_eq!(o.latest, Some(t(2026, 10, 17, 20, 0)));
    assert_eq!(o.saved, 40);
    assert_eq!(src.calls(), vec![today]);
    assert_eq!(
        stored_hours(&db.pool, today).await,
        (0..=20).collect::<Vec<_>>()
    );

    // the 21:00 reading lands; the next poll is current
    let src = FakeRid::new(&[(today, 21)]);
    let o = rid::run_cycle(&src, &db.pool, &stations, t(2026, 10, 17, 21, 40)).await;
    assert!(o.current);
    db.drop().await;
}

#[tokio::test]
async fn yesterdays_late_hours_are_recovered_after_midnight() {
    let Some(db) = TestDb::new().await else {
        return;
    };
    let today = d(2026, 10, 17);
    let yesterday = today - Duration::days(1);
    let stations = Stations::bundled();
    // stored up to 21:00 yesterday (the evening RID published late)
    let partial = FakeRid::new(&[(yesterday, 21)]);
    rid::run_cycle(&partial, &db.pool, &stations, t(2026, 10, 16, 21, 5)).await;
    assert_eq!(
        stored_hours(&db.pool, yesterday).await,
        (1..=21).collect::<Vec<_>>()
    );

    // 02:10: RID now has all of yesterday and today up to 02:00
    let src = FakeRid::new(&[(yesterday, 24), (today, 2)]);
    let o = rid::run_cycle(&src, &db.pool, &stations, t(2026, 10, 17, 2, 10)).await;
    assert!(o.current);
    assert_eq!(src.calls(), vec![yesterday, today]);
    assert_eq!(
        stored_hours(&db.pool, yesterday).await,
        (1..=23).collect::<Vec<_>>()
    );
    assert_eq!(stored_hours(&db.pool, today).await, vec![0, 1, 2]);

    // complete now: the next cycle fetches only today's report
    let src = FakeRid::new(&[(yesterday, 24), (today, 3)]);
    rid::run_cycle(&src, &db.pool, &stations, t(2026, 10, 17, 3, 5)).await;
    assert_eq!(src.calls(), vec![today]);
    db.drop().await;
}

#[tokio::test]
async fn rid_upsert_revises_values_and_rounds_like_numeric() {
    let Some(db) = TestDb::new().await else {
        return;
    };
    let stations = Stations::bundled();
    let r = |level: f64| Reading {
        timestamp: t(2026, 10, 17, 9, 0),
        station_num: 8,
        water_level: level,
        discharge: None,
        discharge_percent: None,
    };
    db::save_rid(&db.pool, &[r(1.0700000524520874)], &stations)
        .await
        .unwrap();
    db::save_rid(&db.pool, &[r(1.0900000333786)], &stations)
        .await
        .unwrap();
    let rows: Vec<(String, f64, Option<f64>, String)> = sqlx::query_as(
        "SELECT s.station_code, w.water_level::float8, w.discharge::float8, w.status
         FROM water_measurements w JOIN stations s ON s.id = w.station_id",
    )
    .fetch_all(&db.pool)
    .await
    .unwrap();
    assert_eq!(rows, vec![("P.1".into(), 1.09, None, "active".into())]);
    db.drop().await;
}

#[tokio::test]
async fn hii_dams_and_rain_round_trip() {
    let Some(db) = TestDb::new().await else {
        return;
    };
    let now = t(2026, 10, 17, 21, 5);
    let rain = RainRecord {
        station_id: 11,
        oldcode: Some("CHM001".into()),
        name_th: None,
        name_en: Some("Chiang Mai 1".into()),
        latitude: Some(18.79),
        longitude: Some(99.0),
        sub_basin_id: Some("601".into()),
        agency: Some("HII".into()),
        timestamp: t(2026, 10, 17, 21, 0),
        rain_1h: Some(1.234),
        rain_24h: Some(12.0),
    };
    assert_eq!(
        db::save_hii_rain(&db.pool, &[rain.clone(), rain.clone()], now)
            .await
            .unwrap(),
        1
    );
    let (r1h,): (f64,) = sqlx::query_as("SELECT rain_1h::float8 FROM hii_rainfall")
        .fetch_one(&db.pool)
        .await
        .unwrap();
    assert_eq!(r1h, 1.23);

    let wl = WaterLevelRecord {
        station_id: 5,
        oldcode: Some("ridhydro_P.1".into()),
        rid_code: Some("P.1".into()),
        name_th: Some("สะพานนวรัฐ".into()),
        name_en: None,
        latitude: None,
        longitude: None,
        agency: None,
        river_name: None,
        offset_msl: Some(300.0),
        ground_level_msl: None,
        min_bank_msl: None,
        critical_level_msl: None,
        critical_level_m: None,
        qmax: None,
        is_key_station: true,
        timestamp: t(2026, 10, 17, 21, 0),
        wl_msl: Some(301.2),
        wl_m: None,
        discharge: None,
        flow_rate: None,
        storage_percent: None,
        situation_level: Some(2),
        diff_wl_bank: None,
    };
    db::save_hii_waterlevel(&db.pool, &[wl], now).await.unwrap();
    let (code, key, sit): (String, bool, i32) = sqlx::query_as(
        "SELECT s.rid_code, s.is_key_station, w.situation_level
         FROM hii_waterlevel w JOIN hii_wl_stations s ON s.id = w.station_id",
    )
    .fetch_one(&db.pool)
    .await
    .unwrap();
    assert_eq!((code.as_str(), key, sit), ("P.1", true, 2));

    let dam = |name: Option<&str>, level: Option<f64>| DamRecord {
        dam_id: "200103".into(),
        region: Some("north".into()),
        name_th: name.map(str::to_string),
        latitude: Some(19.17),
        longitude: None,
        capacity_max_mcm: Some(265.0),
        capacity_normal_mcm: None,
        date: d(2026, 10, 17),
        storage_mcm: Some(200.0),
        storage_pct: Some(75.5),
        inflow_mcm: None,
        outflow_mcm: Some(1.2),
        level_msl: level,
    };
    db::save_dams(&db.pool, &[dam(Some("แม่งัด"), Some(395.1))], now)
        .await
        .unwrap();
    // a later snapshot without name / level must not blank them
    db::save_dams(&db.pool, &[dam(None, None)], now)
        .await
        .unwrap();
    let (name, level): (Option<String>, Option<f64>) = sqlx::query_as(
        "SELECT d.name_th, r.level_msl::float8 FROM rid_reservoir_daily r JOIN rid_dams d USING (dam_id)",
    )
    .fetch_one(&db.pool)
    .await
    .unwrap();
    assert_eq!(name.as_deref(), Some("แม่งัด"));
    assert_eq!(level, Some(395.1));

    let rows = vec![RainRow {
        timestamp: t(2026, 10, 18, 3, 0),
        points: [Some(1.0), None, Some(2.0), None, Some(0.333)],
        catchment_mean: Some(1.111),
    }];
    db::save_rain(&db.pool, &rows).await.unwrap();
    let (mean, missing): (f64, Option<f64>) =
        sqlx::query_as("SELECT catchment_mean::float8, mae_taeng::float8 FROM openmeteo_rain")
            .fetch_one(&db.pool)
            .await
            .unwrap();
    assert_eq!((mean, missing), (1.11, None));
    db.drop().await;
}