poimen-memory/tests/it_pi_source.rs

use mem_ingest::PiSessionSource;
use mem_chunk::RecordSource;
use mem_core::Role;
use futures::stream::StreamExt;
use std::path::PathBuf;

#[tokio::test]
async fn a1_project_from_cwd() {
    let fixture_path = PathBuf::from("fixtures/pi-session-small.jsonl");
    let source = PiSessionSource::new(fixture_path);
    
    let project_key = source.read_project_key().await.expect("Failed to read project key");
    assert_eq!(project_key, "/tmp/my-project", "Project key should come from cwd");
}

#[tokio::test]
async fn a2_role_counts() {
    let fixture_path = PathBuf::from("fixtures/pi-session-small.jsonl");
    let source = PiSessionSource::new(fixture_path);
    let mut stream = source.records();

    let mut user_count = 0;
    let mut assistant_count = 0;
    let mut tool_result_count = 0;
    let mut system_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::ToolResult => tool_result_count += 1,
                Role::System => system_count += 1,
            }
        }
    }

    assert_eq!(user_count, 1, "Should have 1 user message");
    assert_eq!(assistant_count, 2, "Should have 2 assistant messages");
    assert_eq!(tool_result_count, 1, "Should have 1 tool result message");
    assert_eq!(system_count, 1, "Should have 1 system message (compaction)");
}

#[tokio::test]
async fn a3_content_shapes() {
    let fixture_path = PathBuf::from("fixtures/pi-session-small.jsonl");
    let source = PiSessionSource::new(fixture_path);
    let mut stream = source.records();

    let mut has_string_content = false;
    let mut has_block_array_content = false;
    let mut has_structured_content = false;

    while let Some(result) = stream.next().await {
        if let Ok(record) = result {
            // All content should flatten to non-empty text
            assert!(!record.text.is_empty(), "Content should not be empty");

            // Check for different content types by examining the text
            if record.text == "Hello" {
                has_string_content = true;
            } else if record.text.contains("Response text") && record.text.contains("tool_use") {
                has_block_array_content = true;
            } else if record.text.contains("Tool output") {
                has_structured_content = true;
            }
        }
    }

    assert!(has_string_content, "Should have string content");
    assert!(has_block_array_content, "Should have block array content");
    assert!(has_structured_content, "Should have structured content");
}

#[tokio::test]
async fn a4_truncated_tail() {
    // Verify the fixture loads completely without panicking
    let fixture_path = PathBuf::from("fixtures/pi-session-small.jsonl");
    let source = PiSessionSource::new(fixture_path);
    let mut stream = source.records();

    let mut count = 0;
    while let Some(result) = stream.next().await {
        // Malformed lines should be skipped gracefully
        match result {
            Ok(_) => count += 1,
            Err(e) => {
                eprintln!("Skipped line: {}", e);
            }
        }
    }

    // Should have successfully parsed some records
    assert!(count > 0, "Should have parsed at least some records");
}

#[tokio::test]
async fn a5_provenance_includes_record_id() {
    let fixture_path = PathBuf::from("fixtures/pi-session-small.jsonl");
    let source = PiSessionSource::new(fixture_path);
    let mut stream = source.records();

    let mut found_provenance_with_id = false;
    while let Some(result) = stream.next().await {
        if let Ok(record) = result {
            // Provenance should include session ID and record ID
            assert!(record.provenance.source_id.starts_with("pi:"), "Provenance should start with 'pi:'");
            if record.provenance.source_id.contains("msg-") {
                found_provenance_with_id = true;
            }
        }
    }
    assert!(found_provenance_with_id, "Should have found provenance with message ID");
}