Skip to main content

modo/embed/test/
test_helpers.rs

1use std::pin::Pin;
2use std::sync::atomic::{AtomicUsize, Ordering};
3
4use crate::error::Result;
5
6use crate::embed::backend::EmbeddingBackend;
7use crate::embed::convert::to_f32_blob;
8
9/// In-memory embedding backend for unit tests.
10///
11/// Returns a deterministic f32 blob for any input. The blob contains
12/// `dimensions` floats, each set to `0.1`. Tracks call count for assertions.
13pub struct InMemoryBackend {
14    dimensions: usize,
15    call_count: AtomicUsize,
16}
17
18impl InMemoryBackend {
19    /// Create a backend that returns vectors of the given dimensionality.
20    pub fn new(dimensions: usize) -> Self {
21        Self {
22            dimensions,
23            call_count: AtomicUsize::new(0),
24        }
25    }
26
27    /// Number of times [`embed`](EmbeddingBackend::embed) has been called.
28    pub fn call_count(&self) -> usize {
29        self.call_count.load(Ordering::Relaxed)
30    }
31}
32
33impl EmbeddingBackend for InMemoryBackend {
34    fn embed(&self, _input: &str) -> Pin<Box<dyn Future<Output = Result<Vec<u8>>> + Send + '_>> {
35        self.call_count.fetch_add(1, Ordering::Relaxed);
36        let blob = to_f32_blob(&vec![0.1_f32; self.dimensions]);
37        Box::pin(async move { Ok(blob) })
38    }
39
40    fn dimensions(&self) -> usize {
41        self.dimensions
42    }
43
44    fn model_name(&self) -> &str {
45        "test-embedding"
46    }
47}