kcode-k1-audio-classification-projection 0.4.0

use super::*;
use kcode_k1_audio_classification_format::{
    AudioClassificationEventV3, DiscardedV2, FailedV2, LocalSpeakerLabel, PersonId,
    ProgressUpdateV1, ProgressV1, QueueV2, encode_event,
};
use kcode_k1_transaction::{GENESIS_PARENT, SubsystemId, build_signed_transaction};
use tempfile::TempDir;

fn id(value: u8) -> TxId {
    TxId::from_bytes([value; 12])
}

fn queue(fragment_id: FragmentId) -> AudioClassificationEventV3 {
    AudioClassificationEventV3::Queue(QueueV2 {
        audio_object_id: fragment_id,
    })
}

fn apply(
    projection: &Projection,
    callback: TxId,
    event: AudioClassificationEventV3,
) -> AppliedEvent {
    projection
        .apply(callback, &encode_event(&event).unwrap())
        .unwrap()
}

fn canonical(ordering: &K1TxnOrdering) -> TxId {
    let bytes = build_signed_transaction(
        ordering.tip().unwrap_or(GENESIS_PARENT),
        1,
        [1; 32],
        SubsystemId::from_str(SUBSYSTEM).unwrap(),
        b"payload",
        |_| Ok([2; 64]),
    )
    .unwrap();
    let id = TxId::for_transaction(&bytes);
    ordering.submit_txn(&bytes).unwrap();
    id
}

#[test]
fn typed_status_bytes_and_schema_one_rebuild_are_exact() {
    let fragment = id(1);
    let mut status = reduce(None, &queue(fragment), id(2)).unwrap().status;
    let known = PersonId::from_tx_id(id(9));
    status.confirmed_labels = vec![
        SpeakerLabelV1 {
            speaker: LocalSpeakerLabel::new(1).unwrap(),
            person_id: Some(known),
        },
        SpeakerLabelV1 {
            speaker: LocalSpeakerLabel::new(2).unwrap(),
            person_id: None,
        },
    ];
    let bytes = postcard::to_allocvec(&status).unwrap();
    let decoded: FragmentStatus = postcard::from_bytes(&bytes).unwrap();
    assert_eq!(decoded.confirmed_labels, status.confirmed_labels);

    let root = TempDir::new().unwrap();
    let ordering = K1TxnOrdering::open(&root.path().join("ordering")).unwrap();
    let directory = root.path().join("projection");
    let (projection, _) = Projection::open(&directory, &ordering).unwrap();
    apply(&projection, id(3), queue(fragment));
    projection
        .connection
        .lock()
        .unwrap()
        .execute("UPDATE metadata SET schema_version = 1", [])
        .unwrap();
    drop(projection);
    let (projection, cursor) = Projection::open(&directory, &ordering).unwrap();
    assert_eq!(cursor, None);
    assert_eq!(projection.status(fragment).unwrap(), None);
    let version: i64 = projection
        .connection
        .lock()
        .unwrap()
        .query_row("SELECT schema_version FROM metadata", [], |row| row.get(0))
        .unwrap();
    assert_eq!(version, SCHEMA_VERSION);
}

#[test]
fn callback_is_atomic_and_restart_recovers_cursor() {
    let root = TempDir::new().unwrap();
    let ordering = K1TxnOrdering::open(&root.path().join("ordering")).unwrap();
    let directory = root.path().join("projection");
    let (projection, cursor) = Projection::open(&directory, &ordering).unwrap();
    assert_eq!(cursor, None);
    let tail = canonical(&ordering);
    apply(&projection, tail, queue(id(2)));
    drop(projection);
    assert_eq!(
        Projection::open(&directory, &ordering).unwrap().1,
        Some(tail)
    );

    let atomic = root.path().join("atomic");
    let (projection, _) = Projection::open(&atomic, &ordering).unwrap();
    Connection::open(atomic.join(DATABASE_NAME)).unwrap().execute_batch(
        "CREATE TRIGGER block_cursor BEFORE UPDATE ON metadata BEGIN SELECT RAISE(ABORT, 'blocked'); END;",
    ).unwrap();
    assert!(
        projection
            .apply(id(8), &encode_event(&queue(id(8))).unwrap())
            .is_err()
    );
    assert_eq!(projection.status(id(8)).unwrap(), None);
    let cursor: Option<Vec<u8>> = projection
        .connection
        .lock()
        .unwrap()
        .query_row("SELECT last_applied_txid FROM metadata", [], |row| {
            row.get(0)
        })
        .unwrap();
    assert_eq!(cursor, None);
}

#[test]
fn retry_discard_and_independent_roots() {
    let root = TempDir::new().unwrap();
    let ordering = K1TxnOrdering::open(&root.path().join("ordering")).unwrap();
    let (projection, _) = Projection::open(&root.path().join("projection"), &ordering).unwrap();
    let fragment = id(3);
    apply(&projection, id(4), queue(fragment));
    apply(
        &projection,
        id(5),
        AudioClassificationEventV3::Failed(FailedV2 {
            fragment_id: fragment,
            stage: FragmentStageV1::Queue,
            llm_job_sequence: None,
            error: "failed".into(),
        }),
    );
    assert_eq!(
        apply(&projection, id(6), queue(fragment)).effect,
        ProjectionEffect::Start
    );
    assert_eq!(
        projection.status(fragment).unwrap().unwrap().attempt_count,
        0
    );
    apply(
        &projection,
        id(7),
        AudioClassificationEventV3::Progress(ProgressV1 {
            fragment_id: fragment,
            update: ProgressUpdateV1::LlmJobStarted {
                sequence: 1,
                stage: FragmentStageV1::Transcript,
                name: "retry".into(),
            },
        }),
    );
    assert_eq!(
        projection.status(fragment).unwrap().unwrap().attempt_count,
        1
    );
    assert_eq!(
        apply(
            &projection,
            id(8),
            AudioClassificationEventV3::Discarded(DiscardedV2 {
                fragment_id: fragment,
            })
        )
        .effect,
        ProjectionEffect::Abort
    );
    let discarded = projection.status(fragment).unwrap().unwrap();
    assert_eq!(
        apply(&projection, id(9), queue(fragment)).effect,
        ProjectionEffect::None
    );
    assert_eq!(projection.status(fragment).unwrap().unwrap(), discarded);
    let cursor: Vec<u8> = projection
        .connection
        .lock()
        .unwrap()
        .query_row("SELECT last_applied_txid FROM metadata", [], |row| {
            row.get(0)
        })
        .unwrap();
    assert_eq!(cursor, id(9).as_bytes());

    let held = projection.connection.lock().unwrap();
    let (second, _) = Projection::open(&root.path().join("second"), &ordering).unwrap();
    assert!(second.queued().unwrap().is_empty());
    drop(held);
}