tidaldb/tidal/tests/m6_cohort.rs

//! Milestone 6 Phase 1 Integration Tests: Cohort Engine + Cohort-Scoped Trending.
//!
//! Exercises the complete M6P1 scenario end-to-end:
//!
//! 1. Define cohorts (tech, sports) via predicates on user metadata.
//! 2. Write users with metadata placing them in different cohorts.
//! 3. Write items and signals with user context (triggering cohort attribution).
//! 4. Verify cohort signal state is independent per cohort.
//! 5. Execute RETRIEVE with `.cohort("tech")` and verify cohort-scoped ranking.
//! 6. Execute RETRIEVE with `.profile("cohort_trending")` + cohort clause.
//! 7. Verify that signals from non-cohort users do not appear in cohort state.

#![allow(
    clippy::too_many_lines,
    clippy::unwrap_used,
    clippy::cast_precision_loss
)]

use std::{collections::HashMap, time::Duration};

#[cfg(feature = "test-utils")]
use tidaldb::TempTidalHome;
use tidaldb::{
    TidalDb,
    cohort::{CohortDef, Predicate},
    query::retrieve::Retrieve,
    schema::{DecaySpec, EntityId, EntityKind, SchemaBuilder, Timestamp, Window},
};

// ── Schema ──────────────────────────────────────────────────────────────────

fn m6_schema() -> tidaldb::schema::Schema {
    let mut builder = SchemaBuilder::new();

    for &(name, half_life_days) in &[
        ("view", 7),
        ("like", 14),
        ("share", 7),
        ("skip", 1),
        ("completion", 14),
    ] {
        let _ = builder
            .signal(
                name,
                EntityKind::Item,
                DecaySpec::Exponential {
                    half_life: Duration::from_secs(half_life_days * 24 * 3600),
                },
            )
            .windows(&[
                Window::OneHour,
                Window::TwentyFourHours,
                Window::SevenDays,
                Window::AllTime,
            ])
            .velocity(true)
            .add();
    }

    builder.build().expect("m6 schema must be valid")
}

// ── Helpers ─────────────────────────────────────────────────────────────────

fn user_metadata(locale: &str, primary_category: &str) -> HashMap<String, String> {
    let mut m = HashMap::new();
    m.insert("locale".to_string(), locale.to_string());
    m.insert("primary_category".to_string(), primary_category.to_string());
    m
}

fn item_metadata(category: &str, creator_id: u64) -> HashMap<String, String> {
    let mut meta = HashMap::new();
    meta.insert("category".to_string(), category.to_string());
    meta.insert("format".to_string(), "video".to_string());
    meta.insert("creator_id".to_string(), creator_id.to_string());
    meta.insert(
        "created_at".to_string(),
        Timestamp::now().as_nanos().to_string(),
    );
    meta
}

fn open_ephemeral_db() -> TidalDb {
    TidalDb::builder()
        .ephemeral()
        .with_schema(m6_schema())
        .open()
        .expect("db open")
}

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

#[test]
fn define_cohorts_and_verify_registry() {
    let db = open_ephemeral_db();

    // Define two cohorts.
    db.define_cohort(CohortDef {
        name: "tech_en".to_string(),
        predicate: Predicate::And(vec![
            Predicate::Eq {
                field: "locale".into(),
                value: "en".into(),
            },
            Predicate::Eq {
                field: "primary_category".into(),
                value: "tech".into(),
            },
        ]),
    })
    .unwrap();

    db.define_cohort(CohortDef {
        name: "sports".to_string(),
        predicate: Predicate::Eq {
            field: "primary_category".into(),
            value: "sports".into(),
        },
    })
    .unwrap();

    // Duplicate should fail.
    let result = db.define_cohort(CohortDef {
        name: "tech_en".to_string(),
        predicate: Predicate::Eq {
            field: "x".into(),
            value: "y".into(),
        },
    });
    assert!(result.is_err());

    db.close().unwrap();
}

#[test]
fn cohort_signal_attribution_on_signal_with_context() {
    let db = open_ephemeral_db();

    // Define cohorts.
    db.define_cohort(CohortDef {
        name: "tech_en".to_string(),
        predicate: Predicate::And(vec![
            Predicate::Eq {
                field: "locale".into(),
                value: "en".into(),
            },
            Predicate::Eq {
                field: "primary_category".into(),
                value: "tech".into(),
            },
        ]),
    })
    .unwrap();

    db.define_cohort(CohortDef {
        name: "sports".to_string(),
        predicate: Predicate::Eq {
            field: "primary_category".into(),
            value: "sports".into(),
        },
    })
    .unwrap();

    // Write users with metadata.
    let user_1 = EntityId::new(1001);
    let user_2 = EntityId::new(1002);
    let user_3 = EntityId::new(1003);

    db.write_user(user_1, &user_metadata("en", "tech")).unwrap();
    db.write_user(user_2, &user_metadata("en", "sports"))
        .unwrap();
    db.write_user(user_3, &user_metadata("fr", "tech")).unwrap(); // French tech user -- no cohort match for tech_en

    // Write items.
    for i in 1..=5u64 {
        db.write_item_with_metadata(EntityId::new(i), &item_metadata("tech", 100))
            .unwrap();
    }

    let now = Timestamp::now();

    // User 1 (tech_en) views items 1, 2, 3.
    for i in 1..=3u64 {
        db.signal_with_context("view", EntityId::new(i), 1.0, now, Some(1001), Some(100))
            .unwrap();
    }

    // User 2 (sports) views items 2, 4.
    for &i in &[2u64, 4] {
        db.signal_with_context("view", EntityId::new(i), 1.0, now, Some(1002), Some(100))
            .unwrap();
    }

    // User 3 (fr, tech -- not in tech_en cohort) views item 5.
    db.signal_with_context("view", EntityId::new(5), 1.0, now, Some(1003), Some(100))
        .unwrap();

    // Verify cohort signal state.
    let cohort_ledger = db.cohort_ledger();

    // tech_en cohort: items 1, 2, 3 should have signals (from user 1).
    let tech_count_1 = cohort_ledger
        .read_windowed_count("tech_en", EntityId::new(1), "view", Window::AllTime)
        .unwrap();
    assert_eq!(tech_count_1, 1, "tech_en item 1 should have 1 view");

    let tech_count_2 = cohort_ledger
        .read_windowed_count("tech_en", EntityId::new(2), "view", Window::AllTime)
        .unwrap();
    assert_eq!(tech_count_2, 1, "tech_en item 2 should have 1 view");

    // tech_en cohort: item 4 should have 0 signals (only user 2 viewed it).
    let tech_count_4 = cohort_ledger
        .read_windowed_count("tech_en", EntityId::new(4), "view", Window::AllTime)
        .unwrap();
    assert_eq!(
        tech_count_4, 0,
        "tech_en item 4 should have 0 views (sports user)"
    );

    // sports cohort: items 2, 4 should have signals (from user 2).
    let sports_count_2 = cohort_ledger
        .read_windowed_count("sports", EntityId::new(2), "view", Window::AllTime)
        .unwrap();
    assert_eq!(sports_count_2, 1, "sports item 2 should have 1 view");

    let sports_count_4 = cohort_ledger
        .read_windowed_count("sports", EntityId::new(4), "view", Window::AllTime)
        .unwrap();
    assert_eq!(sports_count_4, 1, "sports item 4 should have 1 view");

    // sports cohort: item 1 should have 0 signals (only tech user viewed it).
    let sports_count_1 = cohort_ledger
        .read_windowed_count("sports", EntityId::new(1), "view", Window::AllTime)
        .unwrap();
    assert_eq!(
        sports_count_1, 0,
        "sports item 1 should have 0 views (tech user)"
    );

    // tech_en: item 5 should have 0 (user 3 is fr, not en).
    let tech_count_5 = cohort_ledger
        .read_windowed_count("tech_en", EntityId::new(5), "view", Window::AllTime)
        .unwrap();
    assert_eq!(
        tech_count_5, 0,
        "tech_en item 5 should have 0 views (user 3 is fr locale)"
    );

    db.close().unwrap();
}

#[test]
fn retrieve_with_named_cohort() {
    let db = open_ephemeral_db();

    // Define cohort.
    db.define_cohort(CohortDef {
        name: "tech_en".to_string(),
        predicate: Predicate::And(vec![
            Predicate::Eq {
                field: "locale".into(),
                value: "en".into(),
            },
            Predicate::Eq {
                field: "primary_category".into(),
                value: "tech".into(),
            },
        ]),
    })
    .unwrap();

    // Write user.
    db.write_user(EntityId::new(1001), &user_metadata("en", "tech"))
        .unwrap();

    // Write 5 items.
    for i in 1..=5u64 {
        db.write_item_with_metadata(EntityId::new(i), &item_metadata("tech", 100))
            .unwrap();
    }

    let now = Timestamp::now();

    // User 1001 (tech_en) views items with different weights.
    // Item 1: 10 views, Item 2: 5 views, Item 3: 1 view.
    for _ in 0..10 {
        db.signal_with_context("view", EntityId::new(1), 1.0, now, Some(1001), Some(100))
            .unwrap();
    }
    for _ in 0..5 {
        db.signal_with_context("view", EntityId::new(2), 1.0, now, Some(1001), Some(100))
            .unwrap();
    }
    db.signal_with_context("view", EntityId::new(3), 1.0, now, Some(1001), Some(100))
        .unwrap();

    // Execute RETRIEVE with cohort_trending profile and cohort clause.
    let query = Retrieve::builder()
        .profile("cohort_trending")
        .cohort("tech_en")
        .limit(5)
        .build()
        .unwrap();

    let results = db.retrieve(&query).unwrap();

    // All 5 items should be returned.
    assert!(
        !results.items.is_empty(),
        "cohort query should return items"
    );

    // Item 1 should be ranked highest (most cohort views).
    assert_eq!(
        results.items[0].entity_id,
        EntityId::new(1),
        "item 1 should be ranked first (10 cohort views)"
    );

    db.close().unwrap();
}

#[test]
fn retrieve_with_nonexistent_cohort_fails() {
    let db = open_ephemeral_db();

    // Write at least one item so the query pipeline runs.
    db.write_item_with_metadata(EntityId::new(1), &item_metadata("tech", 100))
        .unwrap();

    let query = Retrieve::builder()
        .profile("cohort_trending")
        .cohort("nonexistent_cohort")
        .limit(5)
        .build()
        .unwrap();

    let result = db.retrieve(&query);
    assert!(result.is_err(), "query with nonexistent cohort should fail");

    db.close().unwrap();
}

#[test]
fn cohort_trending_profile_exists() {
    let db = open_ephemeral_db();

    // Verify the profile is registered and usable.
    let query = Retrieve::builder()
        .profile("cohort_trending")
        .limit(5)
        .build()
        .unwrap();

    // Should succeed (no items, but the profile is found).
    let results = db.retrieve(&query).unwrap();
    assert!(results.items.is_empty());

    db.close().unwrap();
}

#[test]
fn cohort_ledger_decay_score() {
    let db = open_ephemeral_db();

    db.define_cohort(CohortDef {
        name: "test".to_string(),
        predicate: Predicate::Eq {
            field: "primary_category".into(),
            value: "test".into(),
        },
    })
    .unwrap();

    db.write_user(EntityId::new(2001), &user_metadata("en", "test"))
        .unwrap();

    db.write_item_with_metadata(EntityId::new(1), &item_metadata("test", 200))
        .unwrap();

    let now = Timestamp::now();
    db.signal_with_context("view", EntityId::new(1), 5.0, now, Some(2001), Some(200))
        .unwrap();

    let cohort_ledger = db.cohort_ledger();
    let score = cohort_ledger
        .read_decay_score("test", EntityId::new(1), "view", 0)
        .unwrap();
    assert!(score.is_some(), "cohort decay score should be present");
    assert!(
        score.unwrap() > 0.0,
        "cohort decay score should be positive"
    );

    db.close().unwrap();
}

#[test]
fn cohort_ledger_velocity() {
    let db = open_ephemeral_db();

    db.define_cohort(CohortDef {
        name: "test".to_string(),
        predicate: Predicate::Eq {
            field: "primary_category".into(),
            value: "test".into(),
        },
    })
    .unwrap();

    db.write_user(EntityId::new(2001), &user_metadata("en", "test"))
        .unwrap();

    db.write_item_with_metadata(EntityId::new(1), &item_metadata("test", 200))
        .unwrap();

    let now = Timestamp::now();
    // Record 10 signals from the test cohort user.
    for _ in 0..10 {
        db.signal_with_context("view", EntityId::new(1), 1.0, now, Some(2001), Some(200))
            .unwrap();
    }

    let cohort_ledger = db.cohort_ledger();
    let vel = cohort_ledger
        .read_velocity("test", EntityId::new(1), "view", Window::OneHour)
        .unwrap();
    assert!(
        vel > 0.0,
        "cohort velocity should be positive after 10 signals"
    );

    db.close().unwrap();
}

#[test]
fn signal_without_user_context_no_cohort_attribution() {
    let db = open_ephemeral_db();

    db.define_cohort(CohortDef {
        name: "test".to_string(),
        predicate: Predicate::Eq {
            field: "primary_category".into(),
            value: "test".into(),
        },
    })
    .unwrap();

    db.write_item_with_metadata(EntityId::new(1), &item_metadata("test", 200))
        .unwrap();

    let now = Timestamp::now();
    // Signal without user context -- should NOT attribute to any cohort.
    db.signal("view", EntityId::new(1), 1.0, now).unwrap();

    let cohort_ledger = db.cohort_ledger();
    let count = cohort_ledger
        .read_windowed_count("test", EntityId::new(1), "view", Window::AllTime)
        .unwrap();
    assert_eq!(
        count, 0,
        "signals without user context should not be attributed to cohorts"
    );

    db.close().unwrap();
}

#[test]
#[cfg(feature = "test-utils")]
fn cohort_definition_survives_restart() {
    let home = TempTidalHome::new().unwrap();

    // First open: define two cohorts with different predicate shapes.
    {
        let db = TidalDb::builder()
            .with_data_dir(home.path())
            .with_schema(m6_schema())
            .open()
            .unwrap();

        db.define_cohort(CohortDef {
            name: "tech_fans".to_string(),
            predicate: Predicate::And(vec![
                Predicate::Eq {
                    field: "locale".into(),
                    value: "en".into(),
                },
                Predicate::Eq {
                    field: "primary_category".into(),
                    value: "tech".into(),
                },
            ]),
        })
        .unwrap();

        db.define_cohort(CohortDef {
            name: "gamers".to_string(),
            predicate: Predicate::Any {
                field: "interest".into(),
                values: vec!["gaming".into(), "esports".into()],
            },
        })
        .unwrap();

        db.close().unwrap();
    }

    // Second open: cohort definitions must be restored from durable storage.
    {
        let db = TidalDb::builder()
            .with_data_dir(home.path())
            .with_schema(m6_schema())
            .open()
            .unwrap();

        // Verify both cohorts are present.
        // We can exercise define_cohort with the same name to confirm
        // the registry already has them -- it should return Err.
        let dup_result = db.define_cohort(CohortDef {
            name: "tech_fans".to_string(),
            predicate: Predicate::Eq {
                field: "x".into(),
                value: "y".into(),
            },
        });
        assert!(
            dup_result.is_err(),
            "tech_fans should already be registered after restart"
        );

        let dup_result = db.define_cohort(CohortDef {
            name: "gamers".to_string(),
            predicate: Predicate::Eq {
                field: "x".into(),
                value: "y".into(),
            },
        });
        assert!(
            dup_result.is_err(),
            "gamers should already be registered after restart"
        );

        // Verify that a new cohort can still be defined after restart.
        db.define_cohort(CohortDef {
            name: "new_cohort".to_string(),
            predicate: Predicate::Eq {
                field: "locale".into(),
                value: "fr".into(),
            },
        })
        .unwrap();

        db.close().unwrap();
    }

    // Third open: verify the third cohort was also persisted.
    {
        let db = TidalDb::builder()
            .with_data_dir(home.path())
            .with_schema(m6_schema())
            .open()
            .unwrap();

        let dup_result = db.define_cohort(CohortDef {
            name: "new_cohort".to_string(),
            predicate: Predicate::Eq {
                field: "x".into(),
                value: "y".into(),
            },
        });
        assert!(
            dup_result.is_err(),
            "new_cohort should be registered after second restart"
        );

        db.close().unwrap();
    }
}

#[test]
#[cfg(feature = "test-utils")]
fn cohort_definition_survives_restart_and_queries_work() {
    let home = TempTidalHome::new().unwrap();

    // First open: define cohort, write user + items + signals.
    {
        let db = TidalDb::builder()
            .with_data_dir(home.path())
            .with_schema(m6_schema())
            .open()
            .unwrap();

        db.define_cohort(CohortDef {
            name: "tech_en".to_string(),
            predicate: Predicate::And(vec![
                Predicate::Eq {
                    field: "locale".into(),
                    value: "en".into(),
                },
                Predicate::Eq {
                    field: "primary_category".into(),
                    value: "tech".into(),
                },
            ]),
        })
        .unwrap();

        db.write_user(EntityId::new(1001), &user_metadata("en", "tech"))
            .unwrap();

        for i in 1..=3u64 {
            db.write_item_with_metadata(EntityId::new(i), &item_metadata("tech", 100))
                .unwrap();
        }

        let now = Timestamp::now();
        for _ in 0..5 {
            db.signal_with_context("view", EntityId::new(1), 1.0, now, Some(1001), Some(100))
                .unwrap();
        }

        db.close().unwrap();
    }

    // Second open: verify cohort query still works.
    {
        let db = TidalDb::builder()
            .with_data_dir(home.path())
            .with_schema(m6_schema())
            .open()
            .unwrap();

        // The cohort definition should be restored.
        let dup = db.define_cohort(CohortDef {
            name: "tech_en".to_string(),
            predicate: Predicate::Eq {
                field: "x".into(),
                value: "y".into(),
            },
        });
        assert!(
            dup.is_err(),
            "tech_en cohort definition should survive restart"
        );

        // Retrieve with cohort_trending should work (profile resolves the cohort name).
        let query = Retrieve::builder()
            .profile("cohort_trending")
            .cohort("tech_en")
            .limit(5)
            .build()
            .unwrap();

        // This should not error -- the cohort is known.
        let results = db.retrieve(&query).unwrap();
        // Items may or may not have cohort signal state after restart (signals were
        // checkpointed), but the query must not fail due to a missing cohort definition.
        assert!(
            results.items.len() <= 3,
            "at most 3 items should be returned"
        );

        db.close().unwrap();
    }
}

#[test]
fn multiple_cohort_memberships() {
    let db = open_ephemeral_db();

    // Define overlapping cohorts.
    db.define_cohort(CohortDef {
        name: "english".to_string(),
        predicate: Predicate::Eq {
            field: "locale".into(),
            value: "en".into(),
        },
    })
    .unwrap();

    db.define_cohort(CohortDef {
        name: "tech".to_string(),
        predicate: Predicate::Eq {
            field: "primary_category".into(),
            value: "tech".into(),
        },
    })
    .unwrap();

    // User belongs to BOTH cohorts.
    db.write_user(EntityId::new(3001), &user_metadata("en", "tech"))
        .unwrap();

    db.write_item_with_metadata(EntityId::new(1), &item_metadata("tech", 300))
        .unwrap();

    let now = Timestamp::now();
    db.signal_with_context("view", EntityId::new(1), 1.0, now, Some(3001), Some(300))
        .unwrap();

    let cohort_ledger = db.cohort_ledger();

    // Signal should appear in BOTH cohorts.
    let english_count = cohort_ledger
        .read_windowed_count("english", EntityId::new(1), "view", Window::AllTime)
        .unwrap();
    let tech_count = cohort_ledger
        .read_windowed_count("tech", EntityId::new(1), "view", Window::AllTime)
        .unwrap();

    assert_eq!(
        english_count, 1,
        "english cohort should have the signal from en-locale user"
    );
    assert_eq!(
        tech_count, 1,
        "tech cohort should have the signal from tech user"
    );

    db.close().unwrap();
}

/// Regression: defining a cohort at runtime, *after* a user's membership has
/// already been resolved and cached, must invalidate the resolver cache so the
/// next signal for that user is attributed to the newly-defined cohort.
///
/// Before the fix, `CohortResolver` cached per-user memberships keyed by user
/// id and was only ever invalidated per-user on a metadata write. `define_cohort`
/// added the cohort to the registry without touching the cache, so any user
/// already in the cache kept returning the stale (pre-define) membership set:
/// the new cohort silently accumulated zero signal for that user until they were
/// LRU-evicted or their metadata was rewritten. No error, no log.
///
/// Sequence that reproduces it:
///   1. Define cohort A (`tech_en`).
///   2. Write user U and send a signal — this resolves U against {A} and caches
///      the result. (The metadata write invalidates U's entry first, so the
///      signal is what actually populates the cache against the A-only set.)
///   3. Define cohort B (`tech`) whose predicate U also matches.
///   4. Send another signal for U.
///   5. Assert B's ledger received U's signal.
///
/// Without `define_cohort` calling `invalidate_all()`, step 4's resolve returns
/// the cached {A} set, B's ledger stays empty, and the final assert sees a count
/// of 0 instead of 1.
#[test]
fn define_cohort_after_resolve_invalidates_cache_and_attributes() {
    let db = open_ephemeral_db();

    let user = EntityId::new(3001);
    let item = EntityId::new(1);

    // 1. Define cohort A only.
    db.define_cohort(CohortDef {
        name: "tech_en".to_string(),
        predicate: Predicate::And(vec![
            Predicate::Eq {
                field: "locale".into(),
                value: "en".into(),
            },
            Predicate::Eq {
                field: "primary_category".into(),
                value: "tech".into(),
            },
        ]),
    })
    .unwrap();

    // 2. Write user U (en/tech) and item, then signal — this caches U's
    //    membership against the current cohort set {tech_en}. The write_user
    //    invalidates U's entry first; the signal repopulates it.
    db.write_user(user, &user_metadata("en", "tech")).unwrap();
    db.write_item_with_metadata(item, &item_metadata("tech", 100))
        .unwrap();

    let now = Timestamp::now();
    db.signal_with_context("view", item, 1.0, now, Some(user.as_u64()), Some(100))
        .unwrap();

    // Sanity: cohort A saw the first signal.
    let cohort_ledger = db.cohort_ledger();
    assert_eq!(
        cohort_ledger
            .read_windowed_count("tech_en", item, "view", Window::AllTime)
            .unwrap(),
        1,
        "cohort A should have the first signal"
    );

    // 3. Define cohort B at runtime. U also matches B (primary_category=tech).
    //    Without invalidate_all(), U's cached membership still omits B.
    db.define_cohort(CohortDef {
        name: "tech".to_string(),
        predicate: Predicate::Eq {
            field: "primary_category".into(),
            value: "tech".into(),
        },
    })
    .unwrap();

    // 4. Signal U again — must re-resolve against {tech_en, tech} now.
    db.signal_with_context("view", item, 1.0, now, Some(user.as_u64()), Some(100))
        .unwrap();

    // 5. Cohort B must have received the post-define signal. This is the
    //    assertion that fails (count == 0) without the cache invalidation.
    let cohort_ledger = db.cohort_ledger();
    let b_count = cohort_ledger
        .read_windowed_count("tech", item, "view", Window::AllTime)
        .unwrap();
    assert_eq!(
        b_count, 1,
        "cohort B (defined after U was cached) must receive U's post-define signal; \
         a stale resolver cache would leave this at 0"
    );

    // Cohort A should now show both signals — it was always in U's set.
    let a_count = cohort_ledger
        .read_windowed_count("tech_en", item, "view", Window::AllTime)
        .unwrap();
    assert_eq!(a_count, 2, "cohort A should have both signals");

    db.close().unwrap();
}