328 lines
9.9 KiB
Rust
328 lines
9.9 KiB
Rust
#![allow(clippy::unwrap_used)]
|
|
//! PG1 Baseline Comparison — Integration Tests
|
|
//!
|
|
//! End-to-end tests for the A/B experiment framework:
|
|
//! chronological profile, user assignment, metric aggregation, and lift.
|
|
|
|
use std::collections::HashMap;
|
|
use std::time::Duration;
|
|
|
|
use tidaldb::schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window};
|
|
use tidaldb::{ExperimentConfig, ExperimentGroup, TidalDb};
|
|
|
|
/// Build a schema with view, click, and complete 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(
|
|
"click",
|
|
EntityKind::Item,
|
|
DecaySpec::Exponential {
|
|
half_life: Duration::from_secs(7 * 24 * 3600),
|
|
},
|
|
)
|
|
.windows(&[Window::OneHour, Window::TwentyFourHours])
|
|
.velocity(false)
|
|
.add();
|
|
let _ = builder
|
|
.signal(
|
|
"complete",
|
|
EntityKind::Item,
|
|
DecaySpec::Exponential {
|
|
half_life: Duration::from_secs(7 * 24 * 3600),
|
|
},
|
|
)
|
|
.windows(&[Window::OneHour, Window::TwentyFourHours])
|
|
.velocity(false)
|
|
.add();
|
|
builder.build().unwrap()
|
|
}
|
|
|
|
fn make_config() -> ExperimentConfig {
|
|
ExperimentConfig {
|
|
experiment_id: "pg1-baseline-v1".into(),
|
|
treatment_fraction: 0.5,
|
|
treatment_profile: "for_you".into(),
|
|
control_profile: "chronological".into(),
|
|
click_signals: vec!["click".into()],
|
|
completion_signals: vec!["complete".into()],
|
|
return_window: Duration::from_secs(7 * 86400),
|
|
}
|
|
}
|
|
|
|
// ── Chronological Profile Tests ──────────────────────────────────────────────
|
|
|
|
#[test]
|
|
fn chronological_profile_returns_reverse_created_at() {
|
|
let schema = test_schema();
|
|
let db = TidalDb::builder()
|
|
.ephemeral()
|
|
.with_schema(schema)
|
|
.open()
|
|
.unwrap();
|
|
|
|
let now = Timestamp::now().as_nanos();
|
|
// Write items with distinct created_at values.
|
|
for i in 0..10u64 {
|
|
let mut meta = HashMap::new();
|
|
meta.insert("title".to_string(), format!("Item {i}"));
|
|
meta.insert(
|
|
"created_at".to_string(),
|
|
(now + i * 1_000_000_000).to_string(),
|
|
);
|
|
db.write_item_with_metadata(EntityId::new(i + 1), &meta)
|
|
.unwrap();
|
|
}
|
|
|
|
let results = db
|
|
.retrieve(
|
|
&tidaldb::query::Retrieve::builder()
|
|
.profile("chronological")
|
|
.limit(10)
|
|
.build()
|
|
.unwrap(),
|
|
)
|
|
.unwrap();
|
|
|
|
// Items should be in reverse created_at order (newest first).
|
|
let ids: Vec<u64> = results.items.iter().map(|r| r.entity_id.as_u64()).collect();
|
|
assert_eq!(ids.len(), 10);
|
|
// Newest item (id=10, highest created_at) should be first.
|
|
assert_eq!(ids[0], 10, "newest item should be first");
|
|
assert_eq!(ids[9], 1, "oldest item should be last");
|
|
}
|
|
|
|
#[test]
|
|
fn chronological_produces_same_order_for_different_users() {
|
|
let schema = test_schema();
|
|
let db = TidalDb::builder()
|
|
.ephemeral()
|
|
.with_schema(schema)
|
|
.open()
|
|
.unwrap();
|
|
|
|
let now = Timestamp::now().as_nanos();
|
|
for i in 0..5u64 {
|
|
let mut meta = HashMap::new();
|
|
meta.insert(
|
|
"created_at".to_string(),
|
|
(now + i * 1_000_000_000).to_string(),
|
|
);
|
|
db.write_item_with_metadata(EntityId::new(i + 1), &meta)
|
|
.unwrap();
|
|
}
|
|
|
|
let query = tidaldb::query::Retrieve::builder()
|
|
.profile("chronological")
|
|
.limit(5)
|
|
.build()
|
|
.unwrap();
|
|
|
|
let results1 = db.retrieve(&query).unwrap();
|
|
let results2 = db.retrieve(&query).unwrap();
|
|
|
|
let ids1: Vec<u64> = results1
|
|
.items
|
|
.iter()
|
|
.map(|r| r.entity_id.as_u64())
|
|
.collect();
|
|
let ids2: Vec<u64> = results2
|
|
.items
|
|
.iter()
|
|
.map(|r| r.entity_id.as_u64())
|
|
.collect();
|
|
assert_eq!(ids1, ids2, "chronological should be deterministic");
|
|
}
|
|
|
|
// ── Assignment Tests ─────────────────────────────────────────────────────────
|
|
|
|
#[test]
|
|
fn experiment_group_assignment_is_deterministic() {
|
|
let schema = test_schema();
|
|
let db = TidalDb::builder()
|
|
.ephemeral()
|
|
.with_schema(schema)
|
|
.open()
|
|
.unwrap();
|
|
|
|
let config = make_config();
|
|
let group = db.experiment_group(42, &config).unwrap();
|
|
for _ in 0..100 {
|
|
assert_eq!(db.experiment_group(42, &config).unwrap(), group);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn experiment_profile_returns_correct_name() {
|
|
let schema = test_schema();
|
|
let db = TidalDb::builder()
|
|
.ephemeral()
|
|
.with_schema(schema)
|
|
.open()
|
|
.unwrap();
|
|
|
|
let config = make_config();
|
|
let profile = db.experiment_profile(42, &config).unwrap();
|
|
assert!(
|
|
profile == "for_you" || profile == "chronological",
|
|
"profile should be treatment or control, got: {profile}"
|
|
);
|
|
|
|
let group = db.experiment_group(42, &config).unwrap();
|
|
match group {
|
|
ExperimentGroup::Treatment => assert_eq!(profile, "for_you"),
|
|
ExperimentGroup::Control => assert_eq!(profile, "chronological"),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn experiment_group_balanced_split() {
|
|
let schema = test_schema();
|
|
let db = TidalDb::builder()
|
|
.ephemeral()
|
|
.with_schema(schema)
|
|
.open()
|
|
.unwrap();
|
|
|
|
let config = make_config();
|
|
let treatment_count = (0..10_000u64)
|
|
.filter(|&uid| db.experiment_group(uid, &config).unwrap() == ExperimentGroup::Treatment)
|
|
.count();
|
|
|
|
assert!(
|
|
(4750..=5250).contains(&treatment_count),
|
|
"expected ~5000 treatment users, got {treatment_count}"
|
|
);
|
|
}
|
|
|
|
// ── Report Tests ─────────────────────────────────────────────────────────────
|
|
|
|
#[test]
|
|
fn experiment_report_with_synthetic_signals() {
|
|
let schema = test_schema();
|
|
let db = TidalDb::builder()
|
|
.ephemeral()
|
|
.with_schema(schema)
|
|
.open()
|
|
.unwrap();
|
|
|
|
let now = Timestamp::now();
|
|
|
|
// Write some items.
|
|
for i in 0..10u64 {
|
|
let mut meta = HashMap::new();
|
|
meta.insert("title".to_string(), format!("Item {i}"));
|
|
db.write_item_with_metadata(EntityId::new(i + 1), &meta)
|
|
.unwrap();
|
|
}
|
|
|
|
let config = make_config();
|
|
|
|
// Use 100 users; assign them and record different signal patterns.
|
|
let user_ids: Vec<u64> = (1..=100).collect();
|
|
|
|
// Record signals using signal_with_context for user-aware tracking.
|
|
for &uid in &user_ids {
|
|
let group = db.experiment_group(uid, &config).unwrap();
|
|
|
|
// All users get views.
|
|
for item_id in 1..=5u64 {
|
|
db.signal_with_context("view", EntityId::new(item_id), 1.0, now, Some(uid), None)
|
|
.unwrap();
|
|
}
|
|
|
|
// Treatment users get more clicks and completions.
|
|
match group {
|
|
ExperimentGroup::Treatment => {
|
|
for item_id in 1..=3u64 {
|
|
db.signal_with_context(
|
|
"click",
|
|
EntityId::new(item_id),
|
|
1.0,
|
|
now,
|
|
Some(uid),
|
|
None,
|
|
)
|
|
.unwrap();
|
|
db.signal_with_context(
|
|
"complete",
|
|
EntityId::new(item_id),
|
|
1.0,
|
|
now,
|
|
Some(uid),
|
|
None,
|
|
)
|
|
.unwrap();
|
|
}
|
|
}
|
|
ExperimentGroup::Control => {
|
|
// Control: only 1 click, 0 completions.
|
|
db.signal_with_context("click", EntityId::new(1), 1.0, now, Some(uid), None)
|
|
.unwrap();
|
|
}
|
|
}
|
|
}
|
|
|
|
let report = db.experiment_report(&config, &user_ids).unwrap();
|
|
|
|
// Verify group sizes sum to total.
|
|
assert_eq!(
|
|
report.treatment_users + report.control_users,
|
|
user_ids.len()
|
|
);
|
|
|
|
// Treatment should have higher CTR (3/5 vs 1/5 per user).
|
|
assert!(
|
|
report.treatment_metrics.ctr > report.control_metrics.ctr,
|
|
"treatment CTR ({}) should exceed control CTR ({})",
|
|
report.treatment_metrics.ctr,
|
|
report.control_metrics.ctr
|
|
);
|
|
|
|
// Positive CTR lift.
|
|
assert!(
|
|
report.lift.ctr_lift > 0.0,
|
|
"CTR lift ({}) should be positive",
|
|
report.lift.ctr_lift
|
|
);
|
|
|
|
// Treatment should have higher completion rate.
|
|
assert!(
|
|
report.treatment_metrics.completion_rate >= report.control_metrics.completion_rate,
|
|
"treatment completion ({}) should >= control ({})",
|
|
report.treatment_metrics.completion_rate,
|
|
report.control_metrics.completion_rate
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn experiment_report_is_reproducible() {
|
|
let schema = test_schema();
|
|
let db = TidalDb::builder()
|
|
.ephemeral()
|
|
.with_schema(schema)
|
|
.open()
|
|
.unwrap();
|
|
|
|
let config = make_config();
|
|
let user_ids: Vec<u64> = (1..=50).collect();
|
|
|
|
let report1 = db.experiment_report(&config, &user_ids).unwrap();
|
|
let report2 = db.experiment_report(&config, &user_ids).unwrap();
|
|
|
|
assert_eq!(report1.treatment_users, report2.treatment_users);
|
|
assert_eq!(report1.control_users, report2.control_users);
|
|
assert_eq!(report1.experiment_id, report2.experiment_id);
|
|
}
|