use super::*;
fn open_test_db() -> Connection {
let conn = Connection::open_in_memory().expect("in-memory db");
conn.execute_batch(
"CREATE TABLE memories (
id INTEGER PRIMARY KEY AUTOINCREMENT,
namespace TEXT NOT NULL DEFAULT 'global',
name TEXT NOT NULL,
type TEXT NOT NULL DEFAULT 'note',
description TEXT NOT NULL DEFAULT '',
body TEXT NOT NULL DEFAULT '',
body_hash TEXT NOT NULL DEFAULT '',
session_id TEXT,
source TEXT NOT NULL DEFAULT 'agent',
metadata TEXT NOT NULL DEFAULT '{}',
created_at INTEGER NOT NULL DEFAULT (unixepoch()),
updated_at INTEGER NOT NULL DEFAULT (unixepoch()),
deleted_at INTEGER,
UNIQUE(namespace, name)
);",
)
.expect("schema creation must succeed");
conn
}
fn open_temp_queue() -> (Connection, String) {
let path = format!(
"/tmp/test-enrich-dl-{}-{}.sqlite",
std::process::id(),
fastrand::u64(..)
);
let conn = open_queue_db(&path).expect("queue db must open");
(conn, path)
}
fn insert_pending(conn: &Connection, key: &str) -> i64 {
insert_pending_op(conn, key, "memory", "MemoryBindings")
}
fn insert_pending_op(conn: &Connection, key: &str, item_type: &str, operation: &str) -> i64 {
conn.execute(
"INSERT INTO queue (item_key, item_type, status, operation) VALUES (?1, ?2, 'pending', ?3)",
rusqlite::params![key, item_type, operation],
)
.unwrap();
conn.last_insert_rowid()
}
#[test]
fn classify_provider_error_and_not_found_are_hard() {
assert_eq!(
classify_enrich_outcome(&AppError::ProviderError {
code: "400".into(),
message: "context length exceeded".into(),
}),
crate::retry::AttemptOutcome::HardFailure
);
assert_eq!(
classify_enrich_outcome(&AppError::NotFound("memory 'gone' not found".into())),
crate::retry::AttemptOutcome::HardFailure
);
}
#[test]
fn open_queue_db_migrates_operation_column() {
let (conn, path) = open_temp_queue();
drop(conn);
let conn = open_queue_db(&path).expect("second open is idempotent");
let cols: Vec<String> = {
let mut stmt = conn.prepare("PRAGMA table_info(queue)").unwrap();
stmt.query_map([], |r| r.get::<_, String>(1))
.unwrap()
.collect::<Result<Vec<_>, _>>()
.unwrap()
};
assert!(cols.iter().any(|c| c == "operation"));
assert!(cols.iter().any(|c| c == "memory_id"));
let _ = std::fs::remove_file(&path);
}
#[test]
fn enqueue_candidate_tags_operation_and_memory_id() {
let main = open_test_db();
main.execute(
"INSERT INTO memories (namespace, name, body) VALUES ('global', 'mem-x', 'body')",
[],
)
.unwrap();
let mem_id: i64 = main
.query_row("SELECT id FROM memories WHERE name='mem-x'", [], |r| {
r.get(0)
})
.unwrap();
let (queue, path) = open_temp_queue();
enqueue_candidate(&queue, &main, "global", "mem-x", "memory", "MemoryBindings");
let (op, mid): (String, i64) = queue
.query_row(
"SELECT operation, memory_id FROM queue WHERE item_key='mem-x'",
[],
|r| Ok((r.get(0)?, r.get(1)?)),
)
.unwrap();
assert_eq!(op, "MemoryBindings");
assert_eq!(mid, mem_id);
let _ = std::fs::remove_file(&path);
}
#[test]
fn requeue_dead_resurrects_dead_rows() {
let (conn, path) = open_temp_queue();
conn.execute(
"INSERT INTO queue (item_key, item_type, status, operation, attempt, error, error_class, next_retry_at) \
VALUES ('mem-dead', 'memory', 'dead', 'MemoryBindings', 8, 'boom', 'permanent', datetime('now'))",
[],
)
.unwrap();
let n = conn
.execute(
"UPDATE queue SET status='pending', attempt=0, next_retry_at=NULL, \
error=NULL, error_class=NULL \
WHERE status='dead' AND (operation = ?1 OR operation IS NULL)",
rusqlite::params!["MemoryBindings"],
)
.unwrap();
assert_eq!(n, 1);
let (status, attempt, nra): (String, i64, Option<String>) = conn
.query_row(
"SELECT status, attempt, next_retry_at FROM queue WHERE item_key='mem-dead'",
[],
|r| Ok((r.get(0)?, r.get(1)?, r.get(2)?)),
)
.unwrap();
assert_eq!(status, "pending");
assert_eq!(attempt, 0);
assert!(nra.is_none());
let _ = std::fs::remove_file(&path);
}
#[test]
fn skipped_item_keys_excludes_only_skipped_for_operation() {
let (conn, path) = open_temp_queue();
conn.execute(
"INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-vetoed', 'memory', 'skipped', 'BodyEnrich')",
[],
)
.unwrap();
conn.execute(
"INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-pending', 'memory', 'pending', 'BodyEnrich')",
[],
)
.unwrap();
conn.execute(
"INSERT INTO queue (item_key, item_type, status, operation) VALUES ('mem-other-op', 'memory', 'skipped', 'MemoryBindings')",
[],
)
.unwrap();
let keys = skipped_item_keys(&conn, "BodyEnrich").unwrap();
assert!(
keys.contains("mem-vetoed"),
"vetoed BodyEnrich item must be excluded from scan"
);
assert!(
!keys.contains("mem-pending"),
"pending item is still actionable"
);
assert!(
!keys.contains("mem-other-op"),
"skipped item from another operation must not leak"
);
assert_eq!(keys.len(), 1);
let _ = std::fs::remove_file(&path);
}
#[test]
fn cascade_cleanup_delete_targets_memory_id_and_name() {
let (conn, path) = open_temp_queue();
conn.execute(
"INSERT INTO queue (item_key, item_type, status, memory_id) VALUES ('by-id', 'memory', 'done', 42)",
[],
)
.unwrap();
conn.execute(
"INSERT INTO queue (item_key, item_type, status) VALUES ('by-name', 'memory', 'pending')",
[],
)
.unwrap();
let removed = conn
.execute(
"DELETE FROM queue WHERE memory_id = ?1 OR item_key = ?2",
rusqlite::params![42_i64, "by-name"],
)
.unwrap();
assert_eq!(removed, 2);
let remaining: i64 = conn
.query_row("SELECT COUNT(*) FROM queue", [], |r| r.get(0))
.unwrap();
assert_eq!(remaining, 0);
let _ = std::fs::remove_file(&path);
}
#[test]
fn item_type_for_maps_entity_and_memory() {
assert_eq!(
item_type_for(&EnrichOperation::EntityDescriptions),
"entity"
);
assert_eq!(item_type_for(&EnrichOperation::MemoryBindings), "memory");
assert_eq!(item_type_for(&EnrichOperation::AugmentBindings), "memory");
assert_eq!(item_type_for(&EnrichOperation::BodyExtract), "memory");
assert_eq!(
item_type_for(&EnrichOperation::EntityConnect),
"entity_pair"
);
assert_eq!(
item_type_for(&EnrichOperation::CrossDomainBridges),
"entity_pair"
);
}
#[test]
fn item_type_for_key_honours_reembed_prefixes() {
assert_eq!(item_type_for_key("plain-memory-name", "memory"), "memory");
assert_eq!(
item_type_for_key("entity:tokio-runtime", "memory"),
"entity"
);
assert_eq!(item_type_for_key("chunk:42", "memory"), "chunk");
assert_eq!(item_type_for_key("some-entity", "entity"), "entity");
assert_eq!(item_type_for_key("pair:1:2", "memory"), "entity_pair");
}
#[test]
fn prune_dead_orphans_removes_only_orphan_memory_rows() {
let main = open_test_db();
main.execute(
"INSERT INTO memories (namespace, name, body) VALUES ('global', 'alive', 'b')",
[],
)
.unwrap();
let (queue, path) = open_temp_queue();
queue
.execute(
"INSERT INTO queue (item_key, item_type, status, operation, error_class) \
VALUES ('gone', 'memory', 'dead', 'MemoryBindings', 'permanent')",
[],
)
.unwrap();
queue
.execute(
"INSERT INTO queue (item_key, item_type, status, operation, error_class) \
VALUES ('alive', 'memory', 'dead', 'MemoryBindings', 'permanent')",
[],
)
.unwrap();
queue
.execute(
"INSERT INTO queue (item_key, item_type, status, operation) \
VALUES ('some-entity', 'entity', 'dead', 'EntityDescriptions')",
[],
)
.unwrap();
let pruned = prune_dead_orphans(&queue, &main, "MemoryBindings", "global").unwrap();
assert_eq!(pruned, 1, "only the orphan memory row is pruned");
let remaining: Vec<String> = {
let mut stmt = queue
.prepare("SELECT item_key FROM queue ORDER BY item_key")
.unwrap();
stmt.query_map([], |r| r.get::<_, String>(0))
.unwrap()
.collect::<Result<Vec<_>, _>>()
.unwrap()
};
assert_eq!(remaining, vec!["alive", "some-entity"]);
let _ = std::fs::remove_file(&path);
}
#[test]
fn prune_dead_entity_orphans_removes_only_entity_dead_rows() {
let (queue, path) = open_temp_queue();
queue
.execute(
"INSERT INTO queue (item_key, item_type, status, operation, error_class) \
VALUES ('entity:foo', 'entity', 'dead', 'ReEmbed', 'permanent')",
[],
)
.unwrap();
queue
.execute(
"INSERT INTO queue (item_key, item_type, status, operation, error_class) \
VALUES ('mem-dead', 'memory', 'dead', 'MemoryBindings', 'permanent')",
[],
)
.unwrap();
queue
.execute(
"INSERT INTO queue (item_key, item_type, status, operation) \
VALUES ('entity:bar', 'entity', 'pending', 'ReEmbed')",
[],
)
.unwrap();
let pruned = prune_dead_entity_orphans(&queue, "ReEmbed").unwrap();
assert_eq!(pruned, 1, "only the entity dead row is pruned");
let remaining: Vec<String> = {
let mut stmt = queue
.prepare("SELECT item_key FROM queue ORDER BY item_key")
.unwrap();
stmt.query_map([], |r| r.get::<_, String>(0))
.unwrap()
.collect::<Result<Vec<_>, _>>()
.unwrap()
};
assert_eq!(remaining, vec!["entity:bar", "mem-dead"]);
let _ = std::fs::remove_file(&path);
}
#[test]
fn stale_processing_claim_is_reset_after_threshold() {
let (conn, path) = open_temp_queue();
let id = insert_pending(&conn, "mem-stale");
conn.execute(
"UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) - 7200 WHERE id=?1",
rusqlite::params![id],
)
.unwrap();
let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
assert_eq!(reset, 1, "a stale claim older than the threshold is reset");
let (status, claimed_at): (String, Option<i64>) = conn
.query_row(
"SELECT status, claimed_at FROM queue WHERE id=?1",
rusqlite::params![id],
|r| Ok((r.get(0)?, r.get(1)?)),
)
.unwrap();
assert_eq!(status, "pending");
assert!(claimed_at.is_none(), "claimed_at must be cleared on reset");
let _ = std::fs::remove_file(&path);
}
#[test]
fn fresh_processing_claim_is_preserved() {
let (conn, path) = open_temp_queue();
let id = insert_pending(&conn, "mem-fresh");
conn.execute(
"UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) WHERE id=?1",
rusqlite::params![id],
)
.unwrap();
let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
assert_eq!(reset, 0, "a fresh claim must not be reset");
let status: String = conn
.query_row(
"SELECT status FROM queue WHERE id=?1",
rusqlite::params![id],
|r| r.get(0),
)
.unwrap();
assert_eq!(status, "processing");
let _ = std::fs::remove_file(&path);
}
#[test]
fn enqueue_batch_is_atomic() {
let main = open_test_db();
main.execute(
"INSERT INTO memories (namespace, name, body) VALUES ('global', 'm1', 'b'), ('global', 'm2', 'b'), ('global', 'm3', 'b')",
[],
)
.unwrap();
let (mut queue, path) = open_temp_queue();
let reader = Connection::open(&path).unwrap();
let tx = queue.transaction().unwrap();
let tx_conn: &Connection = &tx;
for key in ["m1", "m2", "m3"] {
enqueue_candidate(tx_conn, &main, "global", key, "memory", "MemoryBindings");
}
let before: i64 = reader
.query_row("SELECT COUNT(*) FROM queue", [], |r| r.get::<_, i64>(0))
.unwrap();
assert_eq!(before, 0, "rows must not be visible before commit");
tx.commit().unwrap();
let after: i64 = reader
.query_row("SELECT COUNT(*) FROM queue", [], |r| r.get::<_, i64>(0))
.unwrap();
assert_eq!(after, 3, "all three rows are visible after commit");
let _ = std::fs::remove_file(&path);
}
#[test]
fn heartbeat_updates_claimed_at() {
let (conn, path) = open_temp_queue();
let id = insert_pending(&conn, "mem-heartbeat");
conn.execute(
"UPDATE queue SET status='processing', claimed_at = CAST(strftime('%s','now') AS INTEGER) - 7200 WHERE id=?1",
rusqlite::params![id],
)
.unwrap();
heartbeat(&conn, id).unwrap();
let claimed_at: Option<i64> = conn
.query_row(
"SELECT claimed_at FROM queue WHERE id=?1",
rusqlite::params![id],
|r| r.get(0),
)
.unwrap();
let claimed_at = claimed_at.expect("claimed_at must be set after heartbeat");
let now: i64 = conn
.query_row("SELECT CAST(strftime('%s','now') AS INTEGER)", [], |r| {
r.get::<_, i64>(0)
})
.unwrap();
assert!(
now - claimed_at <= 5,
"claimed_at must be within 5s of now after heartbeat"
);
let reset = reset_stale_processing_claims(&conn, 1800).unwrap();
assert_eq!(reset, 0, "fresh claim survives the sweep after heartbeat");
let _ = std::fs::remove_file(&path);
}
#[test]
fn dequeue_filters_by_operation_isolation() {
let (conn, path) = open_temp_queue();
let _ed_id = insert_pending_op(&conn, "aeroporto-guarulhos", "entity", "EntityDescriptions");
let mb_id = insert_pending_op(&conn, "yt-memory", "memory", "MemoryBindings");
match dequeue_next_pending(&conn, "MemoryBindings", "").unwrap() {
DequeueOutcome::Claimed(row) => {
assert_eq!(row.id, mb_id);
assert_eq!(row.item_key, "yt-memory");
assert_eq!(row.operation, "MemoryBindings");
}
DequeueOutcome::Empty => panic!("expected MB claim"),
}
assert!(matches!(
dequeue_next_pending(&conn, "MemoryBindings", "").unwrap(),
DequeueOutcome::Empty
));
match dequeue_next_pending(&conn, "EntityDescriptions", "").unwrap() {
DequeueOutcome::Claimed(row) => {
assert_eq!(row.item_key, "aeroporto-guarulhos");
assert_eq!(row.operation, "EntityDescriptions");
}
DequeueOutcome::Empty => panic!("expected ED claim"),
}
let _ = std::fs::remove_file(&path);
}
#[test]
fn dequeue_does_not_claim_pair_keys_under_memory_bindings() {
let (conn, path) = open_temp_queue();
let _pair = insert_pending_op(&conn, "pair:21560:159670", "entity_pair", "EntityConnect");
assert!(matches!(
dequeue_next_pending(&conn, "MemoryBindings", "").unwrap(),
DequeueOutcome::Empty
));
match dequeue_next_pending(&conn, "EntityConnect", "").unwrap() {
DequeueOutcome::Claimed(row) => {
assert_eq!(row.item_key, "pair:21560:159670");
}
DequeueOutcome::Empty => panic!("expected EC claim"),
}
let _ = std::fs::remove_file(&path);
}
#[test]
fn validate_claim_rejects_wrong_type_and_key_shape() {
let pair = ClaimedRow {
id: 1,
item_key: "pair:1:2".into(),
item_type: "entity_pair".into(),
operation: "MemoryBindings".into(),
attempt: 1,
};
match validate_claim(&pair, "MemoryBindings", "memory") {
ClaimCheck::SkipWrongType { reason } => {
assert!(reason.contains("wrong_"), "reason={reason}");
}
other => panic!("expected SkipWrongType, got {other:?}"),
}
let entity = ClaimedRow {
id: 2,
item_key: "aeroporto-guarulhos".into(),
item_type: "entity".into(),
operation: "EntityDescriptions".into(),
attempt: 1,
};
assert_eq!(
validate_claim(&entity, "MemoryBindings", "memory"),
ClaimCheck::RequeueWrongOp
);
assert_eq!(
validate_claim(&entity, "EntityDescriptions", "entity"),
ClaimCheck::Ok
);
assert!(is_non_memory_key_shape("pair:1:2"));
assert!(is_non_memory_key_shape("entity:99"));
assert!(!is_non_memory_key_shape("plain-memory-name"));
}
#[test]
fn legacy_unscoped_rows_are_not_claimed() {
let (conn, path) = open_temp_queue();
conn.execute(
"INSERT INTO queue (item_key, item_type, status, operation) \
VALUES ('orphan-legacy', 'memory', 'pending', 'LegacyUnscoped')",
[],
)
.unwrap();
assert!(matches!(
dequeue_next_pending(&conn, "MemoryBindings", "").unwrap(),
DequeueOutcome::Empty
));
let _ = std::fs::remove_file(&path);
}