nautilus_testkit/
events.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Collection utilities for the data events a client emits onto the live runner channel.
17
18use std::time::Duration;
19
20use nautilus_common::messages::DataEvent;
21use nautilus_core::UUID4;
22
23/// Collects the data events received on `rx` within `timeout`.
24///
25/// While the live runner's thread-local sender keeps the channel open, this waits out the full
26/// window and suits absence checks. Prefer [`collect_data_events_until_response`] whenever the
27/// events end with a correlated response.
28pub async fn drain_data_events(
29    rx: &mut tokio::sync::mpsc::UnboundedReceiver<DataEvent>,
30    timeout: Duration,
31) -> Vec<DataEvent> {
32    let mut events = Vec::new();
33    let deadline = tokio::time::Instant::now() + timeout;
34    while let Ok(Some(event)) = tokio::time::timeout_at(deadline, rx.recv()).await {
35        events.push(event);
36    }
37    events
38}
39
40/// Collects the data events received on `rx` up to and including the response correlated with
41/// `request_id`, then drains any further events available without waiting.
42///
43/// `timeout` bounds only the wait through the correlated response, so a passing run returns when
44/// that response arrives instead of waiting out the window.
45///
46/// # Panics
47///
48/// Panics if the channel closes before the correlated response arrives, or if that response does
49/// not arrive within `timeout`.
50pub async fn collect_data_events_until_response(
51    rx: &mut tokio::sync::mpsc::UnboundedReceiver<DataEvent>,
52    request_id: UUID4,
53    timeout: Duration,
54) -> Vec<DataEvent> {
55    let mut events = Vec::new();
56    tokio::time::timeout(timeout, async {
57        loop {
58            let event = rx.recv().await.expect("data event channel closed");
59            let is_correlated_response = matches!(
60                &event,
61                DataEvent::Response(response) if response.correlation_id() == &request_id
62            );
63            events.push(event);
64
65            if is_correlated_response {
66                break;
67            }
68        }
69    })
70    .await
71    .unwrap_or_else(|_| panic!("timed out waiting for data response {request_id}"));
72
73    while let Ok(event) = rx.try_recv() {
74        events.push(event);
75    }
76
77    events
78}
79
80#[cfg(test)]
81mod tests {
82    use nautilus_common::messages::{DataResponse, data::InstrumentsResponse};
83    use nautilus_core::UnixNanos;
84    use nautilus_model::{
85        identifiers::{ClientId, Venue},
86        instruments::{InstrumentAny, stubs::equity_aapl_itch},
87        stubs::TestDefault,
88    };
89    use rstest::rstest;
90
91    use super::*;
92
93    fn instruments_response(correlation_id: UUID4) -> DataEvent {
94        DataEvent::Response(DataResponse::Instruments(InstrumentsResponse::new(
95            correlation_id,
96            ClientId::test_default(),
97            Venue::test_default(),
98            Vec::new(),
99            None,
100            None,
101            UnixNanos::default(),
102            None,
103        )))
104    }
105
106    fn correlation_ids(events: &[DataEvent]) -> Vec<Option<UUID4>> {
107        events
108            .iter()
109            .map(|event| match event {
110                DataEvent::Response(response) => Some(*response.correlation_id()),
111                _ => None,
112            })
113            .collect()
114    }
115
116    #[rstest]
117    #[case(0)]
118    #[case(3)]
119    #[tokio::test(start_paused = true)]
120    async fn test_drain_data_events_collects_events_queued_before_deadline(#[case] count: usize) {
121        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
122        let request_ids: Vec<UUID4> = (0..count).map(|_| UUID4::new()).collect();
123        for request_id in &request_ids {
124            tx.send(instruments_response(*request_id)).unwrap();
125        }
126
127        let events = drain_data_events(&mut rx, Duration::from_millis(50)).await;
128
129        assert_eq!(
130            correlation_ids(&events),
131            request_ids.iter().copied().map(Some).collect::<Vec<_>>()
132        );
133    }
134
135    #[rstest]
136    #[tokio::test(start_paused = true)]
137    async fn test_drain_data_events_stops_at_the_absolute_deadline() {
138        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
139        let first_id = UUID4::new();
140        let second_id = UUID4::new();
141        let late_id = UUID4::new();
142
143        // The 40ms and 80ms sends land inside the 100ms window, the 160ms send past it
144        tokio::spawn(async move {
145            tokio::time::sleep(Duration::from_millis(40)).await;
146            tx.send(instruments_response(first_id)).unwrap();
147            tokio::time::sleep(Duration::from_millis(40)).await;
148            tx.send(instruments_response(second_id)).unwrap();
149            tokio::time::sleep(Duration::from_millis(80)).await;
150            tx.send(instruments_response(late_id)).unwrap();
151        });
152
153        let start = tokio::time::Instant::now();
154        let events = drain_data_events(&mut rx, Duration::from_millis(100)).await;
155
156        assert_eq!(
157            correlation_ids(&events),
158            vec![Some(first_id), Some(second_id)]
159        );
160        assert_eq!(start.elapsed(), Duration::from_millis(100));
161    }
162
163    #[rstest]
164    #[tokio::test(start_paused = true)]
165    async fn test_drain_data_events_returns_when_channel_closes() {
166        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
167        let request_id = UUID4::new();
168        tx.send(instruments_response(request_id)).unwrap();
169        drop(tx);
170
171        let start = tokio::time::Instant::now();
172        let events = drain_data_events(&mut rx, Duration::from_secs(5)).await;
173
174        assert_eq!(correlation_ids(&events), vec![Some(request_id)]);
175        assert_eq!(start.elapsed(), Duration::ZERO);
176    }
177
178    #[rstest]
179    #[tokio::test(start_paused = true)]
180    async fn test_collect_data_events_until_response_returns_at_correlated_response() {
181        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
182        let other_id = UUID4::new();
183        let request_id = UUID4::new();
184        let trailing_id = UUID4::new();
185        tx.send(DataEvent::Instrument(InstrumentAny::Equity(
186            equity_aapl_itch(),
187        )))
188        .unwrap();
189        tx.send(instruments_response(other_id)).unwrap();
190        tx.send(instruments_response(request_id)).unwrap();
191        tx.send(instruments_response(trailing_id)).unwrap();
192
193        let start = tokio::time::Instant::now();
194        let events =
195            collect_data_events_until_response(&mut rx, request_id, Duration::from_secs(5)).await;
196
197        assert_eq!(
198            correlation_ids(&events),
199            vec![None, Some(other_id), Some(request_id), Some(trailing_id)]
200        );
201        assert_eq!(start.elapsed(), Duration::ZERO);
202    }
203
204    #[rstest]
205    #[tokio::test(start_paused = true)]
206    #[should_panic(expected = "timed out waiting for data response")]
207    async fn test_collect_data_events_until_response_panics_without_correlated_response() {
208        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
209        tx.send(instruments_response(UUID4::new())).unwrap();
210
211        collect_data_events_until_response(&mut rx, UUID4::new(), Duration::from_millis(50)).await;
212    }
213
214    #[rstest]
215    #[tokio::test]
216    #[should_panic(expected = "data event channel closed")]
217    async fn test_collect_data_events_until_response_panics_when_channel_closes() {
218        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<DataEvent>();
219        drop(tx);
220
221        collect_data_events_until_response(&mut rx, UUID4::new(), Duration::from_secs(5)).await;
222    }
223}
nautilus_testkit/events.rs

nautilus_testkit/
events.rs
