stemedb/crates/stemedb-query/tests/battery/helpers.rs

//! Shared test utilities for battery tests.

#![allow(clippy::expect_used)] // Test code uses expect() for clear failure messages

// Re-export commonly used items for tests
pub use ed25519_dalek::{Signer, SigningKey};
pub use rand::rngs::OsRng;
pub use std::sync::Arc;
pub use stemedb_core::serde::serialize;
pub use stemedb_core::signing::compute_content_hash_v2;
pub use stemedb_core::testing::AssertionBuilder;
pub use stemedb_core::types::{
    Assertion, EscalationLevel, EscalationPolicy, LifecycleStage, ObjectValue, ResolutionStatus,
    SignatureEntry, SourceClass,
};
pub use stemedb_ingest::worker::{serialize_assertion, IngestWorker};
pub use stemedb_lens::{
    AsyncLens, LayeredConsensusLens, RecencyLens, SyncLensWrapper, TrustAwareAuthorityLens,
};
pub use stemedb_query::{
    apply_source_class_decay, Materializer, Query, QueryEngine, SkepticResolver,
};
pub use stemedb_storage::{
    key_codec, EscalationStore, GenericEscalationStore, GenericIndexStore, GenericTrustRankStore,
    GenericVoteStore, HybridStore, IndexStore, KVStore,
};
pub use stemedb_wal::Journal;
pub use tempfile::tempdir;
pub use tokio::sync::Mutex;

/// Create a signed assertion with Ed25519 signature (v1), source class, confidence,
/// and arbitrary ObjectValue (text/bool, not just number).
///
/// The signature signs the message `"{subject}:{predicate}"` which matches
/// IngestWorker's v1 verification logic.
pub fn create_signed_assertion_with_source(
    subject: &str,
    predicate: &str,
    object: ObjectValue,
    source_class: SourceClass,
    confidence: f32,
    timestamp: u64,
) -> Assertion {
    let mut csprng = OsRng;
    let signing_key = SigningKey::generate(&mut csprng);
    let verifying_key = signing_key.verifying_key();

    let message = format!("{}:{}", subject, predicate);
    let signature = signing_key.sign(message.as_bytes());

    AssertionBuilder::new()
        .subject(subject)
        .predicate(predicate)
        .object(object)
        .source_class(source_class)
        .confidence(confidence)
        .lifecycle(LifecycleStage::Proposed)
        .timestamp(timestamp)
        .signatures(vec![SignatureEntry {
            agent_id: verifying_key.to_bytes(),
            signature: signature.to_bytes(),
            timestamp,
            version: 1,
        }])
        .build()
}

/// Create a signed assertion with Ed25519 signature (v2 enterprise), source class,
/// confidence, and arbitrary ObjectValue.
///
/// v2 signatures sign the BLAKE3 content hash of the full serialized assertion,
/// protecting ALL fields from tampering (confidence, object, timestamp, etc.).
pub fn create_signed_assertion_v2(
    subject: &str,
    predicate: &str,
    object: ObjectValue,
    source_class: SourceClass,
    confidence: f32,
    timestamp: u64,
) -> Assertion {
    let mut csprng = OsRng;
    let signing_key = SigningKey::generate(&mut csprng);
    let verifying_key = signing_key.verifying_key();

    // Build assertion WITHOUT signatures first (for hash computation)
    let mut assertion = AssertionBuilder::new()
        .subject(subject)
        .predicate(predicate)
        .object(object)
        .source_class(source_class)
        .confidence(confidence)
        .lifecycle(LifecycleStage::Proposed)
        .timestamp(timestamp)
        .signatures(vec![])
        .build();

    // Compute the canonical v2 content hash (subject:predicate:object:source_hash:...).
    // Must use compute_content_hash_v2, NOT blake3::hash(serialize(&assertion)) —
    // the verifier checks the canonical fields hash, not the rkyv serialization hash.
    let content_hash = compute_content_hash_v2(&assertion);

    // Sign the content hash (v2 enterprise format)
    let signature = signing_key.sign(&content_hash);

    // Add signature with version 2
    assertion.signatures = vec![SignatureEntry {
        agent_id: verifying_key.to_bytes(),
        signature: signature.to_bytes(),
        timestamp,
        version: 2,
    }];

    assertion
}

/// Store an assertion directly into H: and SP: keys (bypassing WAL/Ingest).
///
/// Used for unit-style tests that don't need the full pipeline.
pub async fn store_assertion_direct(
    store: &Arc<HybridStore>,
    index_store: &GenericIndexStore<Arc<HybridStore>>,
    assertion: &Assertion,
) {
    let bytes = serialize(assertion).expect("serialize assertion");
    let hash = blake3::hash(&bytes);
    let hash_hex = hash.to_hex().to_string();
    let key = key_codec::assertion_key(&assertion.subject, &hash_hex);
    store.put(&key, &bytes).await.expect("put assertion");

    let assertion_hash: [u8; 32] = *hash.as_bytes();
    index_store
        .add_to_indexes(&assertion.subject, &assertion.predicate, &assertion_hash)
        .await
        .expect("add to indexes");
}