krabka_client_streams/
test_driver.rs

1//! Synchronous, broker-free driver for testing topologies.
2//!
3//! This driver is the analog of the JVM `TopologyTestDriver`. Pipe a typed
4//! input record, then read typed output. The driver loops a record back into
5//! the graph when the record goes to an internal repartition topic that is also
6//! a source.
7
8use std::collections::{HashMap, HashSet, VecDeque};
9
10use crate::{
11    embedded::EmbeddedTask,
12    processor::{
13        erased::{OutputRecord, ProcessorError},
14        serde::{Consumed, Produced, Serde, SerdeAssociate},
15    },
16    topology::BuiltTopology,
17};
18
19/// A pending record entry in the repartition loop-back queue.
20type PendingRecord = (String, Option<Vec<u8>>, Vec<u8>, i64);
21
22/// Synchronous test driver.
23///
24/// This driver builds a runnable processor graph from a [`BuiltTopology`]. It
25/// exposes `pipe_input` and `read_output`, so a test can exercise processing
26/// logic without a real broker.
27pub struct TopologyTestDriver {
28    task: EmbeddedTask,
29    source_topics: HashSet<String>,
30    output: HashMap<String, VecDeque<OutputRecord>>,
31    /// Mock wall clock in ms. `advance_wall_clock_time` advances it, and the
32    /// driver passes it to `Graph::punctuate_wall_clock`. It starts at `0`, the
33    /// same as the JVM `TopologyTestDriver` mock time at construction.
34    mock_wall_ms: i64,
35    /// The next offset to stamp on a record of each topic. The driver has no
36    /// broker, so it numbers each topic's records from `0` as they are piped.
37    /// The number reaches a processor through `record_context()`.
38    offsets: HashMap<String, i64>,
39    /// Changelog records the stores produced, collected in order as
40    /// `(changelog_topic, key, value, record_ts)`. The driver has no broker, so
41    /// it keeps the drained changelog here instead of discarding it. A test can
42    /// then assert byte-exact changelog records with `drain_changelog`.
43    changelog_captured: Vec<(String, bytes::Bytes, Option<bytes::Bytes>, Option<i64>)>,
44}
45
46impl TopologyTestDriver {
47    /// Instantiates the topology's graph for testing.
48    ///
49    /// # Errors
50    ///
51    /// Returns the [`ProcessorError`] that instantiation or a processor's
52    /// `init` raises. A topology that `build` accepted raises none in
53    /// practice.
54    pub fn new(built: &BuiltTopology) -> Result<Self, ProcessorError> {
55        let source_topics: HashSet<String> = built.list_source_topics().into_iter().collect();
56        // The driver derives the default changelog names from the application
57        // id `app`, whatever id the topology was built with, and it keeps every
58        // changelog record, reuse-source ones included, so a test can assert
59        // them byte for byte. The task opens in-memory stores with the record
60        // cache disabled, the JVM `TopologyTestDriver` default, so stores emit
61        // on every update and goldens stay deterministic.
62        let task = EmbeddedTask::instantiate(built, "app", HashSet::new(), None)?;
63        Ok(Self {
64            task,
65            source_topics,
66            output: HashMap::new(),
67            mock_wall_ms: 0,
68            offsets: HashMap::new(),
69            changelog_captured: Vec::new(),
70        })
71    }
72
73    /// Serializes and pipes one record on `topic`, then loops repartition
74    /// outputs back.
75    ///
76    /// `consumed` is the same `Consumed::with(key_serde, value_serde)` pair the
77    /// source for `topic` reads with.
78    pub fn pipe_input<KS, VS>(
79        &mut self,
80        topic: &str,
81        consumed: impl Into<Consumed<KS, VS>>,
82        key: Option<KS::Target>,
83        value: VS::Target,
84        timestamp: i64,
85    ) where
86        KS: SerdeAssociate + Serde<KS::Target>,
87        VS: SerdeAssociate + Serde<VS::Target>,
88    {
89        let consumed = consumed.into();
90        let kb = key.as_ref().map(|k| consumed.key_serde.serialize(topic, k));
91        let vb = consumed.value_serde.serialize(topic, &value);
92        drop(key);
93        drop(value);
94        self.pipe_bytes(topic, kb.as_deref(), &vb, timestamp);
95    }
96
97    /// Claims the next offset of `topic`.
98    fn next_offset(&mut self, topic: &str) -> i64 {
99        let offset = self.offsets.entry(topic.to_string()).or_insert(0);
100        let claimed = *offset;
101        *offset += 1;
102        claimed
103    }
104
105    fn pipe_bytes(&mut self, topic: &str, key: Option<&[u8]>, value: &[u8], timestamp: i64) {
106        let mut queue: VecDeque<PendingRecord> = VecDeque::from([(
107            topic.to_string(),
108            key.map(<[u8]>::to_vec),
109            value.to_vec(),
110            timestamp,
111        )]);
112        while let Some((t, k, v, ts)) = queue.pop_front() {
113            // run the graph for this topic; ignore unknown topics
114            let offset = self.next_offset(&t);
115            let _ = self.task.pipe(&t, 0, offset, k.as_deref(), &v, ts);
116            self.route_outputs(&mut queue);
117            // Stream-time advanced by this record → fire stream-time punctuators; their
118            // forwarded records route like any output (and may loop back to a source).
119            let _ = self.task.punctuate_stream_time(self.task.stream_time());
120            self.route_outputs(&mut queue);
121        }
122    }
123
124    /// Advances the mock wall clock by `by` and fires wall-clock punctuators.
125    ///
126    /// Each punctuator fires at most once, with the new clock as its value.
127    /// This mirrors the JVM `TopologyTestDriver.advanceWallClockTime`.
128    pub fn advance_wall_clock_time(&mut self, by: std::time::Duration) {
129        self.mock_wall_ms += i64::try_from(by.as_millis()).unwrap_or(i64::MAX);
130        let mut queue: VecDeque<PendingRecord> = VecDeque::new();
131        let _ = self.task.punctuate_wall_clock(self.mock_wall_ms);
132        self.route_outputs(&mut queue);
133        // Drain any loopback the punctuators produced (and fire stream-time for those).
134        while let Some((t, k, v, ts)) = queue.pop_front() {
135            let offset = self.next_offset(&t);
136            let _ = self.task.pipe(&t, 0, offset, k.as_deref(), &v, ts);
137            self.route_outputs(&mut queue);
138            let _ = self.task.punctuate_stream_time(self.task.stream_time());
139            self.route_outputs(&mut queue);
140        }
141    }
142
143    /// Drains the graph's output buffer.
144    ///
145    /// Outputs to a source topic go back on the queue as a repartition
146    /// loop-back. All other outputs append to `self.output`. This method also
147    /// drains the changelog buffers and keeps them in `changelog_captured`. The
148    /// test driver has no broker, so restore does nothing with fresh stores,
149    /// but a test can assert the changelog bytes with `drain_changelog`. Both
150    /// the record-processing phase and the stream-time-punctuation phase of
151    /// `pipe_bytes` call this method.
152    fn route_outputs(&mut self, queue: &mut VecDeque<PendingRecord>) {
153        for out in self.task.take_output() {
154            if self.source_topics.contains(&out.topic) {
155                // internal repartition topic feeding another subtopology → loop back
156                let vv = out.value.clone().unwrap_or_default().to_vec();
157                queue.push_back((
158                    out.topic.clone(),
159                    out.key.as_ref().map(|b| b.to_vec()),
160                    vv,
161                    out.timestamp,
162                ));
163            } else {
164                self.output
165                    .entry(out.topic.clone())
166                    .or_default()
167                    .push_back(out);
168            }
169        }
170        // Drain changelog buffers so they don't grow unbounded; the test driver
171        // has no broker, so instead of re-producing them we RETAIN them for
172        // byte-exact changelog assertions (`drain_changelog`). No reuse-source
173        // suppression — the driver never re-produces, so no write-back loop.
174        self.changelog_captured.extend(
175            self.task
176                .drain_changelogs()
177                .into_iter()
178                .map(|r| (r.topic, r.key, r.value, r.timestamp)),
179        );
180    }
181
182    /// Inspects a state store's contents after piping.
183    ///
184    /// This method mirrors the JVM `TopologyTestDriver.getKeyValueStore`. It is
185    /// test-only. Do not use it for interactive queries.
186    pub fn get_key_value_store<K: Send + Sync + 'static, V: Send + 'static>(
187        &mut self,
188        name: &str,
189    ) -> Option<&mut dyn crate::store::api::KeyValueStore<K, V>> {
190        self.task.graph_mut().stores.get_kv::<K, V>(name)
191    }
192
193    /// Test-only synchronous store read that calls `block_on` on the async
194    /// store.
195    ///
196    /// The store API is async. This method drives one typed `get` to
197    /// completion, so sync `#[test]` assertions can inspect materialized state.
198    pub fn store_get<K: Send + Sync + 'static, V: Send + 'static>(
199        &mut self,
200        store: &str,
201        key: &K,
202    ) -> Option<V> {
203        let s = self.task.graph_mut().stores.get_kv::<K, V>(store)?;
204        pollster::block_on(s.get(key))
205    }
206
207    /// Latest live value of a key in a versioned store. This is a test helper.
208    pub fn store_get_versioned<K: Send + Sync + 'static, V: Send + 'static>(
209        &mut self,
210        store_name: &str,
211        key: &K,
212    ) -> Option<V> {
213        let store = self
214            .task
215            .graph_mut()
216            .stores
217            .get_versioned::<K, V>(store_name)?;
218        pollster::block_on(store.get(key)).map(|r| r.value)
219    }
220
221    /// Drains the changelog records the stores have produced so far. This is a
222    /// test helper.
223    ///
224    /// The records come back as `(changelog_topic, key, value, record_ts)` in
225    /// production order. A test can then assert byte-exact changelog records,
226    /// for example the KIP-889 versioned changelog, which holds bare value
227    /// bytes and the version timestamp in the record-timestamp field.
228    pub fn drain_changelog(
229        &mut self,
230    ) -> Vec<(String, bytes::Bytes, Option<bytes::Bytes>, Option<i64>)> {
231        std::mem::take(&mut self.changelog_captured)
232    }
233
234    /// Populates a GLOBAL store directly. This method is test-only.
235    ///
236    /// In a real app the global consumer reads all partitions of the source
237    /// topic and applies them to the fully-replicated global store. The test
238    /// driver instead injects values straight into the shared global store
239    /// manager, so a stream-globaltable join can look them up.
240    pub fn pipe_global<K, V>(&mut self, store_name: &str, key: K, value: V)
241    where
242        K: Send + Sync + 'static,
243        V: Send + 'static,
244    {
245        pollster::block_on(self.task.graph().globals.put(store_name, key, value));
246    }
247
248    // ── Interactive-query reads ─────────────────────────────────────────────
249    //
250    // These read the driver's local stores through the *byte* IQ layer
251    // (`StoreRegistry::iq_get` → `&dyn IqQueryable`), then deserialize — the same
252    // path the supervisor serves `KafkaStreams::*_store` queries from. They
253    // exercise the full DSL → store → IQ-byte-read round-trip, so a Rust IQ test
254    // can assert read semantics without standing up the async supervisor.
255
256    /// Interactive-query KV read: the value for `key`, else `None`. The result
257    /// is `None` if the store is absent or is not a key-value store.
258    /// # Panics
259    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
260    pub async fn iq_kv_get<K: 'static, V: 'static>(
261        &self,
262        store: &str,
263        key: &K,
264        ks: &dyn Serde<K>,
265        vs: &dyn Serde<V>,
266    ) -> Option<V> {
267        let q = self.task.graph().stores.iq_get(store)?;
268        let kb = ks.serialize(store, key);
269        let vb = q.iq_kv_get(&kb).await?;
270        Some(vs.deserialize(store, &vb).expect("iq deserialize"))
271    }
272
273    /// Interactive-query KV range read over the inclusive `[lo, hi]` key span,
274    /// in store (memcmp) order. The result is empty if the store is absent or
275    /// is not a KV store.
276    /// # Panics
277    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
278    pub async fn iq_kv_range<K: 'static, V: 'static>(
279        &self,
280        store: &str,
281        lo: &K,
282        hi: &K,
283        ks: &dyn Serde<K>,
284        vs: &dyn Serde<V>,
285    ) -> Vec<(K, V)> {
286        let Some(q) = self.task.graph().stores.iq_get(store) else {
287            return Vec::new();
288        };
289        let lob = ks.serialize(store, lo);
290        let hib = ks.serialize(store, hi);
291        q.iq_kv_range(&lob, &hib)
292            .await
293            .into_iter()
294            .map(|(k, v)| {
295                (
296                    ks.deserialize(store, &k).expect("iq deserialize"),
297                    vs.deserialize(store, &v).expect("iq deserialize"),
298                )
299            })
300            .collect()
301    }
302
303    /// Interactive-query read of every entry in a KV store, in store order.
304    /// # Panics
305    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
306    pub async fn iq_kv_all<K: 'static, V: 'static>(
307        &self,
308        store: &str,
309        ks: &dyn Serde<K>,
310        vs: &dyn Serde<V>,
311    ) -> Vec<(K, V)> {
312        let Some(q) = self.task.graph().stores.iq_get(store) else {
313            return Vec::new();
314        };
315        q.iq_kv_all()
316            .await
317            .into_iter()
318            .map(|(k, v)| {
319                (
320                    ks.deserialize(store, &k).expect("iq deserialize"),
321                    vs.deserialize(store, &v).expect("iq deserialize"),
322                )
323            })
324            .collect()
325    }
326
327    /// Interactive-query approximate entry count for a KV store. The result is
328    /// `0` if the store is absent.
329    pub async fn iq_kv_count(&self, store: &str) -> u64 {
330        match self.task.graph().stores.iq_get(store) {
331            Some(q) => q.iq_kv_approx_count().await,
332            None => 0,
333        }
334    }
335
336    /// Interactive-query window read: `(windowStart, value)` for every window
337    /// of `key` whose start is in the inclusive `[from, to]` span, ascending by
338    /// window start. The result is empty if the store is absent or is not a
339    /// window store.
340    /// # Panics
341    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
342    pub async fn iq_window_fetch<K: 'static, V: 'static>(
343        &self,
344        store: &str,
345        key: &K,
346        from: i64,
347        to: i64,
348        ks: &dyn Serde<K>,
349        vs: &dyn Serde<V>,
350    ) -> Vec<(i64, V)> {
351        let Some(q) = self.task.graph().stores.iq_get(store) else {
352            return Vec::new();
353        };
354        let kb = ks.serialize(store, key);
355        q.iq_window_fetch(&kb, from, to)
356            .await
357            .into_iter()
358            .map(|(start, v)| (start, vs.deserialize(store, &v).expect("iq deserialize")))
359            .collect()
360    }
361
362    /// Interactive-query single-window read: value of the window of `key` that
363    /// starts exactly at `window_start`, else `None`.
364    /// # Panics
365    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
366    pub async fn iq_window_fetch_single<K: 'static, V: 'static>(
367        &self,
368        store: &str,
369        key: &K,
370        window_start: i64,
371        ks: &dyn Serde<K>,
372        vs: &dyn Serde<V>,
373    ) -> Option<V> {
374        let q = self.task.graph().stores.iq_get(store)?;
375        let kb = ks.serialize(store, key);
376        let vb = q.iq_window_fetch_single(&kb, window_start).await?;
377        Some(vs.deserialize(store, &vb).expect("iq deserialize"))
378    }
379
380    /// Runs an `IQv2` query against the single test graph, partition 0.
381    ///
382    /// `Position` is empty, because the driver does not track source offsets.
383    /// Use unit tests on `Position::dominates` for bound behavior.
384    pub async fn query<Q: crate::runtime::iqv2::Query>(
385        &self,
386        req: crate::runtime::iqv2::StateQuery<Q>,
387    ) -> crate::runtime::iqv2::StateQueryResult<Q::Result> {
388        use std::collections::BTreeMap;
389
390        use crate::runtime::iqv2::{
391            request::Position,
392            result::{FailureReason, QueryResult, StateQueryResult},
393        };
394
395        let store_name = req.store.clone();
396        let kind = req.query.store_kind();
397        let Some(store) = self.task.graph().stores.iq_get(&store_name) else {
398            // The driver has full single-graph knowledge, so an unknown store
399            // name is genuinely `DoesNotExist`. (The distributed runtime cannot
400            // distinguish absent-from-topology vs not-on-this-partition, so
401            // `serve_iq2` returns an empty partition map instead — see §9.)
402            let mut m = BTreeMap::new();
403            m.insert(
404                0,
405                QueryResult::Failure {
406                    reason: FailureReason::DoesNotExist,
407                    message: format!("no state store named {store_name:?}"),
408                },
409            );
410            return StateQueryResult::new(m);
411        };
412        if store.kind() != kind {
413            // Mirror `serve_iq2`: a name that resolves to a different store kind
414            // is reported `NotPresent`, not `DoesNotExist`.
415            let mut m = BTreeMap::new();
416            m.insert(
417                0,
418                QueryResult::Failure {
419                    reason: FailureReason::NotPresent,
420                    message: "wrong store kind".into(),
421                },
422            );
423            return StateQueryResult::new(m);
424        }
425        let lowered = req.query.lower();
426        let qr = match store.iq2_execute(&lowered).await {
427            Ok(boxed) => match boxed.downcast::<Q::Result>() {
428                Ok(r) => QueryResult::Success {
429                    result: *r,
430                    position: Position::default(),
431                },
432                Err(_) => QueryResult::Failure {
433                    reason: FailureReason::StoreException,
434                    message: "IQv2 result type mismatch".into(),
435                },
436            },
437            Err(crate::store::iq::Iq2Failure::UnknownQueryType) => QueryResult::Failure {
438                reason: FailureReason::UnknownQueryType,
439                message: "unknown query type".into(),
440            },
441            Err(crate::store::iq::Iq2Failure::KeyTypeMismatch) => QueryResult::Failure {
442                reason: FailureReason::StoreException,
443                message: "key type mismatch".into(),
444            },
445        };
446        let mut m = BTreeMap::new();
447        m.insert(0, qr);
448        StateQueryResult::new(m)
449    }
450
451    /// Interactive-query session read: `((start, end), value)` for every
452    /// session of `key`, in store order. The result is empty if the store is
453    /// absent or is not a session store.
454    /// # Panics
455    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
456    pub async fn iq_session_fetch<K: 'static, V: 'static>(
457        &self,
458        store: &str,
459        key: &K,
460        ks: &dyn Serde<K>,
461        vs: &dyn Serde<V>,
462    ) -> Vec<((i64, i64), V)> {
463        let Some(q) = self.task.graph().stores.iq_get(store) else {
464            return Vec::new();
465        };
466        let kb = ks.serialize(store, key);
467        q.iq_session_fetch_key(&kb)
468            .await
469            .into_iter()
470            .map(|(win, v)| (win, vs.deserialize(store, &v).expect("iq deserialize")))
471            .collect()
472    }
473
474    /// Pops and deserializes the next output record for `topic`.
475    ///
476    /// `produced` is the same `Produced::with(key_serde, value_serde)` pair the
477    /// sink writing `topic` produced with.
478    /// # Panics
479    /// Panics if synchronized client state is poisoned or a response violates an invariant established by protocol validation.
480    pub fn read_output<KS, VS>(
481        &mut self,
482        topic: &str,
483        produced: impl Into<Produced<KS, VS>>,
484    ) -> Option<(Option<KS::Target>, VS::Target)>
485    where
486        KS: SerdeAssociate + Serde<KS::Target>,
487        VS: SerdeAssociate + Serde<VS::Target>,
488    {
489        let produced = produced.into();
490        let out = self.output.get_mut(topic)?.pop_front()?;
491        let key = out.key.map(|b| {
492            produced
493                .key_serde
494                .deserialize(topic, &b)
495                .expect("test: deserialize output key")
496        });
497        let value = produced
498            .value_serde
499            .deserialize(topic, &out.value.unwrap_or_default())
500            .expect("test: deserialize output value");
501        Some((key, value))
502    }
503}
504
505#[cfg(test)]
506mod tests {
507    use assert2::check;
508    use async_trait::async_trait;
509    use krabka_units::prelude::*;
510
511    use super::*;
512    use crate::{
513        processor::{
514            api::{Processor, ProcessorContext},
515            record::Record,
516            serde::StringSerde,
517        },
518        topology::{NodeHandle, Topology},
519    };
520
521    struct Upper;
522    #[async_trait]
523    impl Processor<String, String, String, String> for Upper {
524        async fn process(
525            &mut self,
526            ctx: &mut ProcessorContext<'_, '_, String, String>,
527            r: Record<String, String>,
528        ) {
529            ctx.forward(Record::new(r.key, r.value.to_uppercase(), r.timestamp));
530        }
531    }
532    struct DropEmpty;
533    #[async_trait]
534    impl Processor<String, String, String, String> for DropEmpty {
535        async fn process(
536            &mut self,
537            ctx: &mut ProcessorContext<'_, '_, String, String>,
538            r: Record<String, String>,
539        ) {
540            if !r.value.is_empty() {
541                ctx.forward(r);
542            }
543        }
544    }
545    struct Identity;
546    #[async_trait]
547    impl Processor<String, String, String, String> for Identity {
548        async fn process(
549            &mut self,
550            ctx: &mut ProcessorContext<'_, '_, String, String>,
551            r: Record<String, String>,
552        ) {
553            ctx.forward(r);
554        }
555    }
556
557    fn map_filter() -> crate::topology::BuiltTopology {
558        let mut t = Topology::new();
559        let src: NodeHandle<String, String> = t.add_source("src", ["in"]);
560        let up = t.add_processor("up", || Upper, [&src]);
561        let flt = t.add_processor("flt", || DropEmpty, [&up]);
562        t.add_sink("out", "out", [&flt]);
563        t.build("app").unwrap()
564    }
565
566    #[test]
567    fn test_driver_disables_cache() {
568        // The JVM `TopologyTestDriver` default is `statestore.cache.max.bytes = 0`,
569        // so the driver builds its graph with caching disabled (emit-on-update).
570        let built = map_filter();
571        let d = TopologyTestDriver::new(&built).unwrap();
572        check!(d.task.graph().cache_max_bytes() == ByteSize::ZERO);
573    }
574
575    #[test]
576    fn map_filter_through() {
577        let built = map_filter();
578        let mut d = TopologyTestDriver::new(&built).unwrap();
579        d.pipe_input(
580            "in",
581            Consumed::with(StringSerde, StringSerde),
582            Some("k".to_string()),
583            "hello".to_string(),
584            0,
585        );
586        check!(
587            d.read_output("out", Produced::with(StringSerde, StringSerde))
588                == Some((Some("k".to_string()), "HELLO".to_string()))
589        );
590        d.pipe_input(
591            "in",
592            Consumed::with(StringSerde, StringSerde),
593            Some("k2".to_string()),
594            String::new(),
595            1,
596        );
597        check!(
598            d.read_output("out", Produced::with(StringSerde, StringSerde))
599                == None::<(Option<String>, String)>
600        );
601    }
602
603    #[test]
604    fn repartition_loops_through() {
605        // src(in) -> identity -> sink(rp, internal repartition) ; src(rp) -> up -> sink(out)
606        let mut t = Topology::new();
607        t.add_repartition_topic("rp");
608        let s1: NodeHandle<String, String> = t.add_source("s1", ["in"]);
609        let id = t.add_processor("id", || Identity, [&s1]);
610        t.add_sink("to_rp", "rp", [&id]);
611        let s2: NodeHandle<String, String> = t.add_source("s2", ["rp"]);
612        let up = t.add_processor("up", || Upper, [&s2]);
613        t.add_sink("out", "out", [&up]);
614        let built = t.build("app").unwrap();
615
616        let mut d = TopologyTestDriver::new(&built).unwrap();
617        d.pipe_input(
618            "in",
619            Consumed::with(StringSerde, StringSerde),
620            None,
621            "hi".to_string(),
622            0,
623        );
624        check!(
625            d.read_output("out", Produced::with(StringSerde, StringSerde))
626                == Some((None, "HI".to_string()))
627        );
628    }
629
630    #[test]
631    fn branch_to_two_sinks() {
632        let mut t = Topology::new();
633        let src: NodeHandle<String, String> = t.add_source("src", ["in"]);
634        let up = t.add_processor("up", || Upper, [&src]);
635        t.add_sink("a", "out-a", [&up]);
636        t.add_sink("b", "out-b", [&up]);
637        let built = t.build("app").unwrap();
638        let mut d = TopologyTestDriver::new(&built).unwrap();
639        d.pipe_input(
640            "in",
641            Consumed::with(StringSerde, StringSerde),
642            None,
643            "x".to_string(),
644            0,
645        );
646        check!(
647            d.read_output("out-a", Produced::with(StringSerde, StringSerde))
648                == Some((None, "X".to_string()))
649        );
650        check!(
651            d.read_output("out-b", Produced::with(StringSerde, StringSerde))
652                == Some((None, "X".to_string()))
653        );
654    }
655
656    #[test]
657    fn stateful_count_and_store_inspection() {
658        use crate::processor::serde::I64Serde;
659        struct Counter;
660        #[async_trait]
661        impl Processor<String, String, String, i64> for Counter {
662            async fn process(
663                &mut self,
664                ctx: &mut ProcessorContext<'_, '_, String, i64>,
665                r: Record<String, String>,
666            ) {
667                let n = {
668                    let s = ctx.get_state_store::<String, i64>("counts").unwrap();
669                    let n = s.get(&r.value).await.unwrap_or(0) + 1;
670                    s.put(r.value.clone(), n).await;
671                    n
672                };
673                ctx.forward(Record::new(Some(r.value), n, r.timestamp));
674            }
675        }
676        let mut t = Topology::new();
677        let src: NodeHandle<String, String> = t.add_source("src", ["in"]);
678        let c = t.add_processor("c", || Counter, [&src]);
679        t.add_state_store("counts", StringSerde, I64Serde, [c.name()]);
680        t.add_sink("out", "out", [&c]);
681        let mut d = TopologyTestDriver::new(&t.build("app").unwrap()).unwrap();
682        d.pipe_input(
683            "in",
684            Consumed::with(StringSerde, StringSerde),
685            None,
686            "a".to_string(),
687            0,
688        );
689        d.pipe_input(
690            "in",
691            Consumed::with(StringSerde, StringSerde),
692            None,
693            "a".to_string(),
694            1,
695        );
696        check!(
697            d.read_output("out", Produced::with(StringSerde, I64Serde))
698                == Some((Some("a".to_string()), 1))
699        );
700        check!(
701            d.read_output("out", Produced::with(StringSerde, I64Serde))
702                == Some((Some("a".to_string()), 2))
703        );
704        check!(d.store_get::<String, i64>("counts", &"a".to_string()) == Some(2));
705    }
706
707    #[tokio::test]
708    async fn iqv2_query_kv_via_driver() {
709        use crate::{
710            processor::serde::{Consumed, StringSerde},
711            runtime::iqv2::{KeyQuery, StateQueryRequest},
712        };
713
714        let b = crate::StreamsBuilder::new();
715        b.stream::<String, String>(["in"])
716            .group_by_key()
717            .count("counts");
718        let built = b.build("app").unwrap();
719        let mut d = TopologyTestDriver::new(&built).unwrap();
720        for v in ["a", "a", "b"] {
721            d.pipe_input(
722                "in",
723                Consumed::with(StringSerde, StringSerde),
724                Some(v.to_string()),
725                v.to_string(),
726                0,
727            );
728        }
729
730        let res = d
731            .query(
732                StateQueryRequest::in_store("counts")
733                    .with_query(KeyQuery::<String, i64>::with_key("a".into())),
734            )
735            .await;
736        let only = res.only_partition_result().unwrap();
737        assert_eq!(only.result(), Some(&Some(2)));
738    }
739}
krabka_client_streams/test_driver.rs

krabka_client_streams/
test_driver.rs
