105 lines
4.2 KiB
Rust
105 lines
4.2 KiB
Rust
|
use mem_ingest::{PiSessionSource, ClaudeTranscriptSource};
|
||
|
|
use mem_chunk::RecordSource;
|
||
|
|
use mem_core::Role;
|
||
|
|
use futures::stream::StreamExt;
|
||
|
|
use std::path::PathBuf;
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn a1_project_from_cwd_field() {
|
||
|
|
let fixture_path = PathBuf::from("fixtures/claude-transcript-small.jsonl");
|
||
|
|
let source = ClaudeTranscriptSource::new(fixture_path);
|
||
|
|
|
||
|
|
let project_key = source.read_project_key().await.expect("Failed to read project key");
|
||
|
|
// The fixture has cwd as /tmp/my-project
|
||
|
|
assert_eq!(project_key, "/tmp/my-project", "Project key should come from cwd field");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn a2_same_project_across_sources() {
|
||
|
|
// Both pi and claude fixtures should resolve to the same project
|
||
|
|
let pi_path = PathBuf::from("fixtures/pi-session-small.jsonl");
|
||
|
|
let claude_path = PathBuf::from("fixtures/claude-transcript-small.jsonl");
|
||
|
|
|
||
|
|
let pi_source = PiSessionSource::new(pi_path);
|
||
|
|
let claude_source = ClaudeTranscriptSource::new(claude_path);
|
||
|
|
|
||
|
|
let pi_key = pi_source.read_project_key().await.expect("Failed to read pi project");
|
||
|
|
let claude_key = claude_source.read_project_key().await.expect("Failed to read claude project");
|
||
|
|
|
||
|
|
assert_eq!(pi_key, claude_key, "Both sources should resolve to the same project");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn a3_role_mapping() {
|
||
|
|
let fixture_path = PathBuf::from("fixtures/claude-transcript-small.jsonl");
|
||
|
|
let source = ClaudeTranscriptSource::new(fixture_path);
|
||
|
|
let mut stream = source.records();
|
||
|
|
|
||
|
|
let mut user_count = 0;
|
||
|
|
let mut assistant_count = 0;
|
||
|
|
let mut system_count = 0;
|
||
|
|
let mut ignored_count = 0;
|
||
|
|
|
||
|
|
while let Some(result) = stream.next().await {
|
||
|
|
if let Ok(record) = result {
|
||
|
|
match record.role {
|
||
|
|
Role::User => user_count += 1,
|
||
|
|
Role::Assistant => assistant_count += 1,
|
||
|
|
Role::System => system_count += 1,
|
||
|
|
Role::ToolResult => ignored_count += 1,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
assert_eq!(user_count, 1, "Should have 1 user message");
|
||
|
|
assert_eq!(assistant_count, 2, "Should have 2 assistant messages");
|
||
|
|
assert_eq!(system_count, 1, "Should have 1 system message (api_error)");
|
||
|
|
assert_eq!(ignored_count, 0, "Should have no tool result messages");
|
||
|
|
|
||
|
|
// Verify that attachment, queue-operation, file-history-snapshot, summary, mode, etc are skipped
|
||
|
|
// Total records should be user + assistant + system = 4
|
||
|
|
assert_eq!(user_count + assistant_count + system_count, 4, "Only relevant types should be emitted");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn a4_shared_flattener() {
|
||
|
|
// Verify that both sources handle content flattening correctly
|
||
|
|
// The key is that when both sources encounter the same content shapes,
|
||
|
|
// they flatten them identically using the shared flatten_content function
|
||
|
|
let pi_path = PathBuf::from("fixtures/pi-session-small.jsonl");
|
||
|
|
let claude_path = PathBuf::from("fixtures/claude-transcript-small.jsonl");
|
||
|
|
|
||
|
|
let pi_source = PiSessionSource::new(pi_path);
|
||
|
|
let claude_source = ClaudeTranscriptSource::new(claude_path);
|
||
|
|
|
||
|
|
let mut pi_stream = pi_source.records();
|
||
|
|
let mut claude_stream = claude_source.records();
|
||
|
|
|
||
|
|
let mut pi_texts = Vec::new();
|
||
|
|
let mut claude_texts = Vec::new();
|
||
|
|
|
||
|
|
while let Some(result) = pi_stream.next().await {
|
||
|
|
if let Ok(record) = result {
|
||
|
|
pi_texts.push(record.text);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
while let Some(result) = claude_stream.next().await {
|
||
|
|
if let Ok(record) = result {
|
||
|
|
claude_texts.push(record.text);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Both should have produced records
|
||
|
|
assert!(!pi_texts.is_empty(), "pi source should produce texts");
|
||
|
|
assert!(!claude_texts.is_empty(), "claude source should produce texts");
|
||
|
|
|
||
|
|
// Both should handle array content (with tool_use blocks)
|
||
|
|
assert!(pi_texts.iter().any(|t| t.contains("[tool_use")), "pi should flatten array content with tool_use");
|
||
|
|
assert!(claude_texts.iter().any(|t| t.contains("[tool_use")), "claude should flatten array content with tool_use");
|
||
|
|
|
||
|
|
// Both should handle simple string content
|
||
|
|
assert!(pi_texts.iter().any(|t| t == "Hello"), "pi should have simple string content");
|
||
|
|
assert!(claude_texts.iter().any(|t| t == "Hello Claude"), "claude should have simple string content");
|
||
|
|
}
|