tidaldb/tidal/tests/pg1_instrumented_metrics.rs

#![allow(clippy::unwrap_used)]
//! PG1 Instrumented Metrics Pipeline — Integration Tests
//!
//! End-to-end tests verifying per-signal-type counters, latency percentiles,
//! personalization staleness, feedback-loop latency, and /diagnostics endpoint.
//!
//! Requires `--features metrics,test-utils`.

use std::collections::HashMap;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::time::Duration;

use tidaldb::TidalDb;
use tidaldb::schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window};

/// Build a schema with view and like signals.
fn test_schema() -> tidaldb::schema::Schema {
    let mut builder = SchemaBuilder::new();
    let _ = builder
        .signal(
            "view",
            EntityKind::Item,
            DecaySpec::Exponential {
                half_life: Duration::from_secs(7 * 24 * 3600),
            },
        )
        .windows(&[Window::OneHour, Window::TwentyFourHours])
        .velocity(false)
        .add();
    let _ = builder
        .signal(
            "like",
            EntityKind::Item,
            DecaySpec::Exponential {
                half_life: Duration::from_secs(7 * 24 * 3600),
            },
        )
        .windows(&[Window::OneHour])
        .velocity(false)
        .add();
    builder.build().unwrap()
}

/// Make an HTTP GET request to the given address and path.
fn http_get(addr: std::net::SocketAddr, path: &str) -> (u16, String) {
    let mut stream = TcpStream::connect(addr).expect("connect");
    write!(
        stream,
        "GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"
    )
    .unwrap();
    stream.flush().unwrap();
    let mut response = String::new();
    stream.read_to_string(&mut response).unwrap();
    let status: u16 = response
        .lines()
        .next()
        .and_then(|l| l.split_whitespace().nth(1))
        .and_then(|s| s.parse().ok())
        .unwrap_or(0);
    let body = response
        .split_once("\r\n\r\n")
        .map_or_else(String::new, |(_, b)| b.to_string());
    (status, body)
}

fn open_db_with_metrics() -> (std::net::SocketAddr, TidalDb) {
    let db = TidalDb::builder()
        .ephemeral()
        .with_schema(test_schema())
        .enable_metrics("127.0.0.1:0")
        .open()
        .expect("open should succeed");
    let addr = db
        .metrics_addr()
        .expect("metrics_addr should be Some when metrics enabled");
    (addr, db)
}

// ── Tests ──────────────────────────────────────────────────────────────────

#[test]
fn prometheus_contains_per_type_counters() {
    let (addr, db) = open_db_with_metrics();

    // Write some signals of different types.
    for i in 1..=10 {
        db.signal("view", EntityId::new(i), 1.0, Timestamp::now())
            .unwrap();
    }
    for i in 1..=5 {
        db.signal("like", EntityId::new(i), 1.0, Timestamp::now())
            .unwrap();
    }

    let (status, body) = http_get(addr, "/metrics");
    assert_eq!(status, 200);

    // Verify per-type counters.
    assert!(
        body.contains("tidaldb_signal_writes_by_type{signal_type=\"view\"}"),
        "should contain view per-type counter"
    );
    assert!(
        body.contains("tidaldb_signal_writes_by_type{signal_type=\"like\"}"),
        "should contain like per-type counter"
    );

    db.close().unwrap();
}

#[test]
fn prometheus_contains_percentile_gauges_after_signals() {
    let (addr, db) = open_db_with_metrics();

    // Write enough signals to populate the latency histogram.
    for i in 1..=20 {
        db.signal("view", EntityId::new(i), 1.0, Timestamp::now())
            .unwrap();
    }

    let (status, body) = http_get(addr, "/metrics");
    assert_eq!(status, 200);

    // After writing signals, the signal_write_latency histogram should have data,
    // so percentile gauges should appear.
    assert!(
        body.contains("tidaldb_signal_write_latency_us_p50"),
        "should contain signal write p50 gauge"
    );
    assert!(
        body.contains("tidaldb_signal_write_latency_us_p95"),
        "should contain signal write p95 gauge"
    );
    assert!(
        body.contains("tidaldb_signal_write_latency_us_p99"),
        "should contain signal write p99 gauge"
    );

    db.close().unwrap();
}

#[test]
fn prometheus_contains_staleness_histogram_after_user_query() {
    let (addr, db) = open_db_with_metrics();

    // Write item metadata so retrieve has something to find.
    let mut meta = HashMap::new();
    meta.insert("title".to_string(), "Test Item".to_string());
    meta.insert("category".to_string(), "tech".to_string());
    db.write_item_with_metadata(EntityId::new(1), &meta)
        .unwrap();

    // Write a signal with user context.
    db.signal_with_context(
        "view",
        EntityId::new(1),
        1.0,
        Timestamp::now(),
        Some(42),
        None,
    )
    .unwrap();

    // Run a retrieve query for that user.
    let query = tidaldb::query::retrieve::Retrieve::builder()
        .for_user(42)
        .build()
        .unwrap();
    let _ = db.retrieve(&query);

    let (status, body) = http_get(addr, "/metrics");
    assert_eq!(status, 200);

    // Staleness histogram should be present (even if empty -- the histogram lines always render).
    assert!(
        body.contains("tidaldb_personalization_staleness_us"),
        "should contain staleness histogram"
    );

    db.close().unwrap();
}

#[test]
fn diagnostics_endpoint_returns_valid_json() {
    let (addr, db) = open_db_with_metrics();

    // Write some signals to populate metrics.
    for i in 1..=5 {
        db.signal("view", EntityId::new(i), 1.0, Timestamp::now())
            .unwrap();
    }

    let (status, body) = http_get(addr, "/diagnostics");
    assert_eq!(status, 200, "diagnostics should return 200");

    // Verify it's valid JSON by checking structure.
    assert!(body.starts_with('{'), "should start with {{");
    assert!(body.ends_with('}'), "should end with }}");
    assert!(
        body.contains("\"signal_writes_by_type\""),
        "should contain signal_writes_by_type key"
    );
    assert!(
        body.contains("\"signal_write_latency_us\""),
        "should contain signal_write_latency_us key"
    );
    assert!(
        body.contains("\"retrieve_latency_us\""),
        "should contain retrieve_latency_us key"
    );
    assert!(
        body.contains("\"search_latency_us\""),
        "should contain search_latency_us key"
    );
    assert!(
        body.contains("\"personalization_staleness_us\""),
        "should contain personalization_staleness_us key"
    );
    assert!(
        body.contains("\"feedback_loop_latency_us\""),
        "should contain feedback_loop_latency_us key"
    );

    // Verify view count is present in the per-type breakdown.
    assert!(body.contains("\"view\""), "should contain view signal type");

    db.close().unwrap();
}