Skip to main content

ados_sdk/testing/
mod.rs

1//! Test harness for vision plugins.
2//!
3//! Mirrors the Python `ados.sdk.testing` ergonomics for the Rust SDK: a plugin
4//! author exercises a frame-consuming plugin without a real plugin host, vision
5//! engine, shared memory, or socket. [`FakeVisionEngine`] emits synthetic
6//! frames (from an in-memory list or a directory of raw frame files) into the
7//! same [`FrameCallback`] a plugin registers with
8//! `ctx.vision.subscribe_frames`, and captures the [`DetectionBatch`]es the
9//! plugin would publish so a test can assert against them.
10//!
11//! The fake builds a real frame ring through the
12//! [`framebus`] contract and resolves each synthetic
13//! frame the same way the production client does (per-slot seqlock, latest-wins,
14//! torn/stale reads dropped), so a plugin's frame-handling path is exercised
15//! end to end, only the host and the OS shared-memory are faked.
16//!
17//! ```no_run
18//! use ados_sdk::testing::FakeVisionEngine;
19//! use ados_protocol::framebus::{FrameFormat, Detection, BoundingBox};
20//! use std::sync::{Arc, Mutex};
21//!
22//! # async fn demo() {
23//! let mut engine = FakeVisionEngine::new("uvc-0", 64, 48, FrameFormat::Rgb24);
24//! let seen = Arc::new(Mutex::new(0usize));
25//! let s = seen.clone();
26//! engine.on_frame(move |frame| {
27//!     // a plugin's per-frame handler runs here
28//!     assert_eq!(frame.descriptor.camera_id, "uvc-0");
29//!     *s.lock().unwrap() += 1;
30//! });
31//! engine.push_solid(0x80); // one grey frame
32//! engine.deliver_all();
33//! assert_eq!(*seen.lock().unwrap(), 1);
34//! # }
35//! ```
36
37use std::path::Path;
38use std::sync::{Arc, Mutex};
39
40use ados_protocol::framebus::{
41    self, DetectionBatch, FrameDescriptor, FrameFormat, RingLayout, FRAMEBUS_DESCRIPTOR_VERSION,
42};
43
44use crate::vision::{Frame, FrameCallback};
45
46/// Default slot count for the fake ring. Large enough that the harness never
47/// recycles a slot under a pending read in a single-threaded test.
48const DEFAULT_SLOT_COUNT: u32 = 8;
49
50/// One captured published detection, in publish order.
51pub type CapturedDetection = DetectionBatch;
52
53/// In-process stand-in for the vision engine + plugin host bridge.
54///
55/// Owns a synthetic frame ring and a queue of pending frames. A test registers
56/// a [`FrameCallback`] with [`on_frame`](Self::on_frame), enqueues synthetic
57/// frames, then calls [`deliver_all`](Self::deliver_all) (or
58/// [`deliver_one`](Self::deliver_one)) to drive them through the resolver into
59/// the callback. Detections the plugin publishes are captured via the sink
60/// [`detection_sink`](Self::detection_sink) returns.
61pub struct FakeVisionEngine {
62    camera_id: String,
63    shm_name: String,
64    format: FrameFormat,
65    layout: RingLayout,
66    region: Vec<u8>,
67    /// Monotonic frame sequence, also the ring slot via `seq % slot_count`.
68    next_seq: u64,
69    /// Pixel payloads waiting to be delivered.
70    pending: Vec<Vec<u8>>,
71    callback: Option<FrameCallback>,
72    captured: Arc<Mutex<Vec<CapturedDetection>>>,
73}
74
75impl FakeVisionEngine {
76    /// A fake engine for one camera at a fixed frame size and format. The ring
77    /// is sized to hold full `width` x `height` frames.
78    pub fn new(camera_id: &str, width: u32, height: u32, format: FrameFormat) -> Self {
79        Self::with_slot_count(camera_id, width, height, format, DEFAULT_SLOT_COUNT)
80    }
81
82    /// As [`new`](Self::new) but with an explicit slot count, for tests that
83    /// want to force slot recycling and exercise the torn/stale-read drop path.
84    pub fn with_slot_count(
85        camera_id: &str,
86        width: u32,
87        height: u32,
88        format: FrameFormat,
89        slot_count: u32,
90    ) -> Self {
91        let layout = RingLayout::for_frame(slot_count, width, height, format);
92        let mut region = vec![0u8; layout.total_len()];
93        layout
94            .write_header(&mut region)
95            .expect("freshly sized region holds its own header");
96        Self {
97            camera_id: camera_id.to_string(),
98            shm_name: format!("ados-vision-{camera_id}"),
99            format,
100            layout,
101            region,
102            next_seq: 0,
103            pending: Vec::new(),
104            callback: None,
105            captured: Arc::new(Mutex::new(Vec::new())),
106        }
107    }
108
109    /// The frame size of one full frame in this ring's format.
110    pub fn frame_bytes(&self) -> usize {
111        self.layout.slot_bytes as usize
112    }
113
114    /// Register the per-frame callback the plugin under test would pass to
115    /// `ctx.vision.subscribe_frames`. Replaces any prior callback.
116    pub fn on_frame<F>(&mut self, callback: F)
117    where
118        F: Fn(Frame) + Send + Sync + 'static,
119    {
120        self.callback = Some(Arc::new(callback));
121    }
122
123    /// Register a callback as a pre-built [`FrameCallback`] (e.g. the exact
124    /// `Arc` a plugin handed to the SDK).
125    pub fn on_frame_arc(&mut self, callback: FrameCallback) {
126        self.callback = Some(callback);
127    }
128
129    /// Enqueue a raw pixel frame. The bytes must be at most one full frame
130    /// (`frame_bytes`); a shorter slice is delivered as a partial frame, which
131    /// is what a real engine does for a truncated capture.
132    pub fn push_frame(&mut self, pixels: Vec<u8>) {
133        self.pending.push(pixels);
134    }
135
136    /// Enqueue a frame filled with one byte value (a flat colour), sized to a
137    /// full frame. Handy for asserting the callback sees the right bytes.
138    pub fn push_solid(&mut self, value: u8) {
139        self.pending.push(vec![value; self.frame_bytes()]);
140    }
141
142    /// Enqueue every `*.bin` / `*.raw` file in `dir` (sorted by name) as a
143    /// frame, reading each file's bytes verbatim. Returns the count enqueued.
144    pub fn push_dir(&mut self, dir: impl AsRef<Path>) -> std::io::Result<usize> {
145        let mut paths: Vec<_> = std::fs::read_dir(dir)?
146            .filter_map(Result::ok)
147            .map(|e| e.path())
148            .filter(|p| {
149                matches!(
150                    p.extension().and_then(|s| s.to_str()),
151                    Some("bin") | Some("raw")
152                )
153            })
154            .collect();
155        paths.sort();
156        let mut n = 0;
157        for p in paths {
158            let bytes = std::fs::read(&p)?;
159            self.pending.push(bytes);
160            n += 1;
161        }
162        Ok(n)
163    }
164
165    /// Deliver the next pending frame: write it into the ring, build its
166    /// descriptor, resolve it through the seqlock, and invoke the callback.
167    /// Returns `true` if a frame was delivered (a frame existed and resolved),
168    /// `false` if the queue was empty. A frame that fails to resolve (a forced
169    /// torn read) returns `true` for "consumed" but does not fire the callback.
170    pub fn deliver_one(&mut self) -> bool {
171        if self.pending.is_empty() {
172            return false;
173        }
174        let pixels = self.pending.remove(0);
175        let seq = self.next_seq + 1;
176        self.next_seq = seq;
177        let slot = (seq % self.layout.slot_count as u64) as u32;
178        framebus::write_slot(&mut self.region, &self.layout, slot, seq, &pixels)
179            .expect("synthetic frame fits a full-size slot");
180
181        let descriptor = FrameDescriptor {
182            v: FRAMEBUS_DESCRIPTOR_VERSION,
183            camera_id: self.camera_id.clone(),
184            frame_id: seq,
185            ts_ns: 1_700_000_000_000 + seq,
186            ts_sensor_ns: None,
187            exposure_ns: None,
188            trigger_seq: None,
189            timebase: ados_protocol::framebus::TIMEBASE_HOST,
190            width: 0, // width/height are descriptive; the resolver keys on byte_len.
191            height: 0,
192            format: self.layout_format(),
193            shm_name: self.shm_name.clone(),
194            slot,
195            seq,
196            byte_len: pixels.len() as u32,
197        };
198
199        if let Some(frame) = self.resolve(descriptor) {
200            if let Some(cb) = &self.callback {
201                cb(frame);
202            }
203        }
204        true
205    }
206
207    /// Deliver every pending frame in order. Returns the number delivered.
208    pub fn deliver_all(&mut self) -> usize {
209        let mut n = 0;
210        while self.deliver_one() {
211            n += 1;
212        }
213        n
214    }
215
216    /// The shared detection sink. Pass this into a plugin (or call
217    /// [`capture`](Self::capture) directly) so detections the plugin would
218    /// publish via `ctx.vision.publish_detection` are recorded here.
219    pub fn detection_sink(&self) -> Arc<Mutex<Vec<CapturedDetection>>> {
220        self.captured.clone()
221    }
222
223    /// Record a detection batch the plugin under test published, as if the
224    /// engine received it. A test that cannot route the real publish call can
225    /// invoke this from a stubbed publish path.
226    pub fn capture(&self, batch: DetectionBatch) {
227        self.captured.lock().expect("capture lock").push(batch);
228    }
229
230    /// The captured detections in publish order. Returns a copy.
231    pub fn captured_detections(&self) -> Vec<CapturedDetection> {
232        self.captured.lock().expect("capture lock").clone()
233    }
234
235    /// Clear the captured detections.
236    pub fn clear_captured(&self) {
237        self.captured.lock().expect("capture lock").clear();
238    }
239
240    /// Resolve a descriptor against the in-memory ring (no `/dev/shm`),
241    /// validating the seqlock exactly as the production resolver does.
242    fn resolve(&self, descriptor: FrameDescriptor) -> Option<Frame> {
243        // The frame owns its pixels (handed to the callback by value), so each
244        // resolve fills a fresh buffer sized to the descriptor.
245        let mut pixels = Vec::with_capacity(descriptor.byte_len as usize);
246        framebus::read_slot_into(
247            &self.region,
248            &self.layout,
249            descriptor.slot,
250            descriptor.seq,
251            &mut pixels,
252        )
253        .ok()
254        .flatten()?;
255        Some(Frame { descriptor, pixels })
256    }
257
258    /// The format the ring was sized for. The ring header records only
259    /// `slot_bytes`, so the harness keeps the format it was built with.
260    fn layout_format(&self) -> FrameFormat {
261        self.format
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use ados_protocol::framebus::{BoundingBox, Detection};
269
270    #[test]
271    fn delivers_a_solid_frame_to_the_callback() {
272        let mut engine = FakeVisionEngine::new("uvc-0", 4, 4, FrameFormat::Rgb24);
273        let seen: Arc<Mutex<Vec<Frame>>> = Arc::new(Mutex::new(Vec::new()));
274        let s = seen.clone();
275        engine.on_frame(move |frame| s.lock().unwrap().push(frame));
276
277        engine.push_solid(0x42);
278        assert_eq!(engine.deliver_all(), 1);
279
280        let got = seen.lock().unwrap();
281        assert_eq!(got.len(), 1);
282        assert_eq!(got[0].descriptor.camera_id, "uvc-0");
283        assert_eq!(got[0].descriptor.frame_id, 1);
284        assert_eq!(got[0].pixels.len(), engine.frame_bytes());
285        assert!(got[0].pixels.iter().all(|&b| b == 0x42));
286    }
287
288    #[test]
289    fn delivers_frames_in_order() {
290        let mut engine = FakeVisionEngine::new("uvc-0", 2, 2, FrameFormat::Rgb24);
291        let ids: Arc<Mutex<Vec<u64>>> = Arc::new(Mutex::new(Vec::new()));
292        let i = ids.clone();
293        engine.on_frame(move |frame| i.lock().unwrap().push(frame.descriptor.frame_id));
294
295        engine.push_solid(1);
296        engine.push_solid(2);
297        engine.push_solid(3);
298        assert_eq!(engine.deliver_all(), 3);
299        assert_eq!(*ids.lock().unwrap(), vec![1, 2, 3]);
300    }
301
302    #[test]
303    fn captures_published_detections() {
304        let engine = FakeVisionEngine::new("uvc-0", 2, 2, FrameFormat::Rgb24);
305        let batch = DetectionBatch {
306            v: ados_protocol::framebus::VISION_DETECTION_VERSION,
307            model_id: "com.example.weeds".into(),
308            camera_id: "uvc-0".into(),
309            frame_id: 1,
310            ts_ms: 1,
311            frame_width: 640,
312            frame_height: 480,
313            detections: vec![Detection {
314                bbox: Some(BoundingBox {
315                    x: 0.0,
316                    y: 0.0,
317                    width: 1.0,
318                    height: 1.0,
319                }),
320                class_label: "weed".into(),
321                confidence: 0.5,
322                track_id: None,
323                assoc_confidence: None,
324                lock_state: None,
325                attributes: None,
326                mask: None,
327                keypoints: None,
328                depth: None,
329                world_pos: None,
330            }],
331        };
332        engine.capture(batch.clone());
333        let got = engine.captured_detections();
334        assert_eq!(got, vec![batch]);
335        engine.clear_captured();
336        assert!(engine.captured_detections().is_empty());
337    }
338
339    #[test]
340    fn push_dir_reads_raw_frames_in_name_order() {
341        let dir = tempfile::tempdir().unwrap();
342        std::fs::write(dir.path().join("b.bin"), [2u8; 12]).unwrap();
343        std::fs::write(dir.path().join("a.raw"), [1u8; 12]).unwrap();
344        std::fs::write(dir.path().join("skip.txt"), b"not a frame").unwrap();
345
346        let mut engine = FakeVisionEngine::new("uvc-0", 2, 2, FrameFormat::Rgb24); // 12-byte frames
347        let n = engine.push_dir(dir.path()).unwrap();
348        assert_eq!(n, 2);
349
350        let first_byte: Arc<Mutex<Vec<u8>>> = Arc::new(Mutex::new(Vec::new()));
351        let f = first_byte.clone();
352        engine.on_frame(move |frame| f.lock().unwrap().push(frame.pixels[0]));
353        engine.deliver_all();
354        // a.raw (value 1) sorts before b.bin (value 2).
355        assert_eq!(*first_byte.lock().unwrap(), vec![1, 2]);
356    }
357}