use tempfile::tempdir;
use collet::agent::session::{SessionSnapshot, SessionStore};
use collet::api::content::Content;
use collet::api::models::Message;
fn make_snapshot(id: &str, completed: bool, timestamp: &str) -> SessionSnapshot {
SessionSnapshot {
session_id: id.to_string(),
working_dir: "/tmp/test".to_string(),
system_prompt: "You are a helpful assistant.".to_string(),
messages: vec![
Message {
role: "user".to_string(),
content: Some(Content::text("Hello")),
reasoning_content: None,
tool_calls: None,
tool_call_id: None,
},
Message {
role: "assistant".to_string(),
content: Some(Content::text("Hi there!")),
reasoning_content: None,
tool_calls: None,
tool_call_id: None,
},
],
last_reasoning: None,
timestamp: timestamp.to_string(),
completed,
user_task: Some("Say hello".to_string()),
model: Some("glm-4.7".to_string()),
ui_state: None,
}
}
#[tokio::test]
async fn save_and_load_round_trip() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let original = make_snapshot("sess-001", false, "2025-01-15T10:00:00Z");
store.save(&original).await.unwrap();
let loaded = store.load("sess-001").await.unwrap();
assert_eq!(loaded.session_id, "sess-001");
assert_eq!(loaded.working_dir, "/tmp/test");
assert_eq!(loaded.system_prompt, "You are a helpful assistant.");
assert_eq!(loaded.messages.len(), 2);
assert_eq!(loaded.messages[0].role, "user");
assert_eq!(
loaded.messages[0]
.content
.as_ref()
.map(|c| c.text_content()),
Some("Hello".to_string())
);
assert_eq!(loaded.messages[1].role, "assistant");
assert_eq!(
loaded.messages[1]
.content
.as_ref()
.map(|c| c.text_content()),
Some("Hi there!".to_string())
);
assert!(!loaded.completed);
assert_eq!(loaded.user_task.as_deref(), Some("Say hello"));
assert_eq!(loaded.model.as_deref(), Some("glm-4.7"));
}
#[tokio::test]
async fn save_preserves_optional_fields() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let mut snap = make_snapshot("sess-opt", true, "2025-02-01T12:00:00Z");
snap.last_reasoning = Some("I thought about it carefully.".to_string());
snap.user_task = None;
snap.model = None;
store.save(&snap).await.unwrap();
let loaded = store.load("sess-opt").await.unwrap();
assert_eq!(
loaded.last_reasoning.as_deref(),
Some("I thought about it carefully.")
);
assert!(loaded.user_task.is_none());
assert!(loaded.model.is_none());
}
#[tokio::test]
async fn find_incomplete_returns_none_when_completed() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let snap = make_snapshot("sess-done", true, "2025-01-15T10:00:00Z");
store.save(&snap).await.unwrap();
let result = store.find_incomplete().await;
assert!(
result.is_none(),
"find_incomplete should return None for a completed session"
);
}
#[tokio::test]
async fn find_incomplete_returns_session_when_not_completed() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let snap = make_snapshot("sess-wip", false, "2025-01-15T10:00:00Z");
store.save(&snap).await.unwrap();
let result = store.find_incomplete().await;
assert!(
result.is_some(),
"find_incomplete should find the incomplete session"
);
assert_eq!(result.unwrap().session_id, "sess-wip");
}
#[tokio::test]
async fn find_incomplete_returns_latest_incomplete() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let completed = make_snapshot("sess-old", true, "2025-01-01T10:00:00Z");
store.save(&completed).await.unwrap();
let incomplete = make_snapshot("sess-new", false, "2025-01-15T10:00:00Z");
store.save(&incomplete).await.unwrap();
let result = store.find_incomplete().await;
assert!(result.is_some());
assert_eq!(result.unwrap().session_id, "sess-new");
}
#[tokio::test]
async fn list_sessions_ordered_by_timestamp_descending() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let snap_a = make_snapshot("sess-a", true, "2025-01-10T10:00:00Z");
let snap_b = make_snapshot("sess-b", false, "2025-01-20T10:00:00Z");
let snap_c = make_snapshot("sess-c", true, "2025-01-15T10:00:00Z");
store.save(&snap_a).await.unwrap();
store.save(&snap_b).await.unwrap();
store.save(&snap_c).await.unwrap();
let sessions = store.list().await;
assert_eq!(sessions.len(), 3);
assert_eq!(sessions[0].0, "sess-b");
assert_eq!(sessions[1].0, "sess-c");
assert_eq!(sessions[2].0, "sess-a");
}
#[tokio::test]
async fn list_empty_store_returns_empty() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let sessions = store.list().await;
assert!(sessions.is_empty());
}
#[tokio::test]
async fn cleanup_removes_old_completed_sessions() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
for i in 1..=4 {
let snap = make_snapshot(
&format!("sess-{i:03}"),
true,
&format!("2025-01-{i:02}T10:00:00Z"),
);
store.save(&snap).await.unwrap();
}
let removed = store.cleanup(2).await.unwrap();
assert_eq!(removed, 2, "should remove 2 old completed sessions");
let remaining = store.list().await;
assert_eq!(remaining.len(), 2);
let ids: Vec<&str> = remaining.iter().map(|(id, _, _)| id.as_str()).collect();
assert!(ids.contains(&"sess-004"), "newest should remain");
assert!(ids.contains(&"sess-003"), "second newest should remain");
}
#[tokio::test]
async fn cleanup_does_not_remove_incomplete_sessions() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let c1 = make_snapshot("completed-1", true, "2025-01-01T10:00:00Z");
let c2 = make_snapshot("completed-2", true, "2025-01-02T10:00:00Z");
let inc = make_snapshot("incomplete-1", false, "2025-01-03T10:00:00Z");
store.save(&c1).await.unwrap();
store.save(&c2).await.unwrap();
store.save(&inc).await.unwrap();
let removed = store.cleanup(1).await.unwrap();
assert_eq!(removed, 1, "should remove 1 old completed session");
let remaining = store.list().await;
let ids: Vec<&str> = remaining.iter().map(|(id, _, _)| id.as_str()).collect();
assert!(
ids.contains(&"incomplete-1"),
"incomplete session should survive cleanup"
);
assert!(
ids.contains(&"completed-2"),
"most recent completed should survive"
);
}
#[tokio::test]
async fn cleanup_with_nothing_to_remove() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let snap = make_snapshot("only-one", true, "2025-01-01T10:00:00Z");
store.save(&snap).await.unwrap();
let removed = store.cleanup(5).await.unwrap();
assert_eq!(removed, 0);
}
#[tokio::test]
async fn load_nonexistent_session_returns_error() {
let dir = tempdir().unwrap();
let store = SessionStore::new(dir.path().to_str().unwrap());
let result = store.load("does-not-exist").await;
assert!(result.is_err());
}