nautilus_common/msgbus/
stubs.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
16use std::{
17    any::Any,
18    cell::RefCell,
19    rc::Rc,
20    sync::{
21        Arc,
22        atomic::{AtomicBool, Ordering},
23    },
24};
25
26use ahash::AHashMap;
27use nautilus_core::UUID4;
28use ustr::Ustr;
29
30use crate::{
31    messages::data::{DataResponse, DataResponseEnvelope},
32    msgbus::{
33        Handler, IntoHandler, ShareableMessageHandler, TypedHandler, TypedIntoHandler,
34        typed_handler::shareable_handler,
35    },
36};
37
38/// Minimal message payload for message bus tests.
39#[derive(Debug, Clone)]
40pub struct StubMessage;
41
42/// No-op handler for message bus tests.
43#[derive(Debug, Clone)]
44pub struct StubMessageHandler {
45    id: Ustr,
46}
47
48impl Handler<dyn Any> for StubMessageHandler {
49    fn id(&self) -> Ustr {
50        self.id
51    }
52
53    fn handle(&self, _message: &dyn Any) {}
54}
55
56#[must_use]
57pub fn get_stub_shareable_handler(id: Option<Ustr>) -> ShareableMessageHandler {
58    let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
59    shareable_handler(Rc::new(StubMessageHandler { id: unique_id }))
60}
61
62/// Handler that tracks whether it has been called.
63#[derive(Debug, Clone)]
64pub struct CallCheckHandler {
65    id: Ustr,
66    called: Arc<AtomicBool>,
67}
68
69impl CallCheckHandler {
70    #[must_use]
71    pub fn new(id: Option<Ustr>) -> Self {
72        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
73        Self {
74            id: unique_id,
75            called: Arc::new(AtomicBool::new(false)),
76        }
77    }
78
79    #[must_use]
80    pub fn was_called(&self) -> bool {
81        self.called.load(Ordering::SeqCst)
82    }
83
84    /// Returns a `ShareableMessageHandler` for registration.
85    #[must_use]
86    pub fn handler(&self) -> ShareableMessageHandler {
87        shareable_handler(Rc::new(self.clone()))
88    }
89}
90
91impl Handler<dyn Any> for CallCheckHandler {
92    fn id(&self) -> Ustr {
93        self.id
94    }
95
96    fn handle(&self, _message: &dyn Any) {
97        self.called.store(true, Ordering::SeqCst);
98    }
99}
100
101/// Creates a call-checking handler and returns both the handler for registration
102/// and a clone that can be used to check if it was called.
103#[must_use]
104pub fn get_call_check_handler(id: Option<Ustr>) -> (ShareableMessageHandler, CallCheckHandler) {
105    let checker = CallCheckHandler::new(id);
106    let handler = checker.handler();
107    (handler, checker)
108}
109
110/// Handler that saves messages it receives (for Any-based routing).
111#[derive(Debug, Clone)]
112pub struct AnySavingHandler<T> {
113    id: Ustr,
114    messages: Rc<RefCell<Vec<T>>>,
115}
116
117impl<T: Clone + 'static> AnySavingHandler<T> {
118    #[must_use]
119    pub fn new(id: Option<Ustr>) -> Self {
120        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
121        Self {
122            id: unique_id,
123            messages: Rc::new(RefCell::new(Vec::new())),
124        }
125    }
126
127    #[must_use]
128    pub fn get_messages(&self) -> Vec<T> {
129        self.messages.borrow().clone()
130    }
131
132    pub fn clear(&self) {
133        self.messages.borrow_mut().clear();
134    }
135
136    /// Returns a `ShareableMessageHandler` for registration.
137    #[must_use]
138    pub fn handler(&self) -> ShareableMessageHandler {
139        shareable_handler(Rc::new(self.clone()))
140    }
141}
142
143impl<T: Clone + 'static> Handler<dyn Any> for AnySavingHandler<T> {
144    fn id(&self) -> Ustr {
145        self.id
146    }
147
148    fn handle(&self, message: &dyn Any) {
149        if let Some(m) = message.downcast_ref::<T>() {
150            self.messages.borrow_mut().push(m.clone());
151        } else {
152            log::error!(
153                "AnySavingHandler: expected {} got {:?}",
154                std::any::type_name::<T>(),
155                message.type_id()
156            );
157        }
158    }
159}
160
161/// Creates an Any-based message saving handler and returns both the handler
162/// for registration and a clone that can be used to retrieve messages.
163#[must_use]
164pub fn get_any_saving_handler<T: Clone + 'static>(
165    id: Option<Ustr>,
166) -> (ShareableMessageHandler, AnySavingHandler<T>) {
167    let saver = AnySavingHandler::new(id);
168    let handler = saver.handler();
169    (handler, saver)
170}
171
172// Type alias for backward compatibility
173pub type MessageSavingHandler<T> = AnySavingHandler<T>;
174
175/// Saves typed payloads from terminal historical envelopes for payload-focused tests.
176#[must_use]
177pub fn get_response_saving_handler<T: Clone + 'static>(
178    id: Option<Ustr>,
179) -> (TypedHandler<DataResponseEnvelope>, AnySavingHandler<T>) {
180    let (payload_handler, saver) = get_any_saving_handler::<T>(id);
181
182    let handler = TypedHandler::from(move |envelope: &DataResponseEnvelope| {
183        if !envelope.completed() {
184            return;
185        }
186
187        let Some(response) = envelope.response() else {
188            return;
189        };
190
191        match response {
192            DataResponse::Data(r) => payload_handler.0.handle(r),
193            DataResponse::Instrument(r) => payload_handler.0.handle(r.as_ref()),
194            DataResponse::Instruments(r) => payload_handler.0.handle(r),
195            DataResponse::Book(r) => payload_handler.0.handle(r),
196            DataResponse::BookDeltas(r) => payload_handler.0.handle(r),
197            DataResponse::BookDepth(r) => payload_handler.0.handle(r),
198            DataResponse::Quotes(r) => payload_handler.0.handle(r),
199            DataResponse::Trades(r) => payload_handler.0.handle(r),
200            DataResponse::FundingRates(r) => payload_handler.0.handle(r),
201            DataResponse::OptionChainReferencePrice(r) => payload_handler.0.handle(r),
202            DataResponse::Bars(r) => payload_handler.0.handle(r),
203        }
204    });
205
206    (handler, saver)
207}
208
209/// Typed handler which saves the messages it receives (no downcast needed).
210#[derive(Debug, Clone)]
211pub struct TypedMessageSavingHandler<T> {
212    id: Ustr,
213    messages: Rc<RefCell<Vec<T>>>,
214}
215
216impl<T: Clone + 'static> TypedMessageSavingHandler<T> {
217    #[must_use]
218    pub fn new(id: Option<Ustr>) -> Self {
219        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
220        Self {
221            id: unique_id,
222            messages: Rc::new(RefCell::new(Vec::new())),
223        }
224    }
225
226    #[must_use]
227    pub fn get_messages(&self) -> Vec<T> {
228        self.messages.borrow().clone()
229    }
230
231    /// Returns a `TypedHandler` that can be used for subscriptions.
232    #[must_use]
233    pub fn handler(&self) -> TypedHandler<T> {
234        TypedHandler::new(self.clone())
235    }
236}
237
238impl<T: Clone + 'static> Handler<T> for TypedMessageSavingHandler<T> {
239    fn id(&self) -> Ustr {
240        self.id
241    }
242
243    fn handle(&self, message: &T) {
244        self.messages.borrow_mut().push(message.clone());
245    }
246}
247
248/// Creates a typed message saving handler and returns both the handler for subscriptions
249/// and a clone that can be used to retrieve messages.
250#[must_use]
251pub fn get_typed_message_saving_handler<T: Clone + 'static>(
252    id: Option<Ustr>,
253) -> (TypedHandler<T>, TypedMessageSavingHandler<T>) {
254    let saving_handler = TypedMessageSavingHandler::new(id);
255    let typed_handler = saving_handler.handler();
256    (typed_handler, saving_handler)
257}
258
259/// Ownership-based typed handler which saves the messages it receives.
260///
261/// Unlike [`TypedMessageSavingHandler`] which borrows messages, this handler
262/// takes ownership which is required for `IntoEndpointMap` endpoints.
263#[derive(Debug, Clone)]
264pub struct TypedIntoMessageSavingHandler<T> {
265    id: Ustr,
266    messages: Rc<RefCell<Vec<T>>>,
267}
268
269impl<T: 'static> TypedIntoMessageSavingHandler<T> {
270    #[must_use]
271    pub fn new(id: Option<Ustr>) -> Self {
272        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
273        Self {
274            id: unique_id,
275            messages: Rc::new(RefCell::new(Vec::new())),
276        }
277    }
278
279    /// Creates a handler backed by an existing shared messages vec.
280    #[must_use]
281    pub fn new_with_messages(id: Option<Ustr>, messages: Rc<RefCell<Vec<T>>>) -> Self {
282        let unique_id = id.unwrap_or_else(|| Ustr::from(UUID4::new().as_str()));
283        Self {
284            id: unique_id,
285            messages,
286        }
287    }
288
289    #[must_use]
290    pub fn get_messages(&self) -> Vec<T>
291    where
292        T: Clone,
293    {
294        self.messages.borrow().clone()
295    }
296
297    /// Returns a `TypedIntoHandler` that can be used for endpoint registration.
298    #[must_use]
299    pub fn handler(&self) -> TypedIntoHandler<T> {
300        TypedIntoHandler::new(Self {
301            id: self.id,
302            messages: self.messages.clone(),
303        })
304    }
305
306    pub fn clear(&self) {
307        self.messages.borrow_mut().clear();
308    }
309}
310
311impl<T: 'static> IntoHandler<T> for TypedIntoMessageSavingHandler<T> {
312    fn id(&self) -> Ustr {
313        self.id
314    }
315
316    fn handle(&self, message: T) {
317        self.messages.borrow_mut().push(message);
318    }
319}
320
321/// Creates an ownership-based typed message saving handler and returns both the handler
322/// for endpoint registration and a clone that can be used to retrieve messages.
323#[must_use]
324pub fn get_typed_into_message_saving_handler<T: 'static>(
325    id: Option<Ustr>,
326) -> (TypedIntoHandler<T>, TypedIntoMessageSavingHandler<T>) {
327    let saving_handler = TypedIntoMessageSavingHandler::new(id);
328    let typed_handler = saving_handler.handler();
329    (typed_handler, saving_handler)
330}
331
332// Legacy API for tests that use the old pattern with thread_local storage.
333// These wrap AnySavingHandler in thread_local for simpler test usage.
334
335thread_local! {
336    static SAVING_HANDLERS: RefCell<AHashMap<Ustr, Box<dyn std::any::Any>>> = RefCell::new(AHashMap::new());
337}
338
339/// Creates a message saving handler and stores it for later retrieval.
340#[must_use]
341pub fn get_message_saving_handler<T: Clone + 'static>(id: Option<Ustr>) -> ShareableMessageHandler {
342    let (handler, saver) = get_any_saving_handler::<T>(id);
343    let handler_id = handler.0.id();
344    SAVING_HANDLERS.with(|handlers| {
345        handlers.borrow_mut().insert(handler_id, Box::new(saver));
346    });
347    handler
348}
349
350/// Retrieves saved messages from a handler created by `get_message_saving_handler`.
351#[must_use]
352pub fn get_saved_messages<T: Clone + 'static>(handler: &ShareableMessageHandler) -> Vec<T> {
353    let handler_id = handler.0.id();
354    SAVING_HANDLERS.with(|handlers| {
355        let handlers = handlers.borrow();
356        if let Some(saver) = handlers.get(&handler_id)
357            && let Some(saver) = saver.downcast_ref::<AnySavingHandler<T>>()
358        {
359            return saver.get_messages();
360        }
361        Vec::new()
362    })
363}
nautilus_common/msgbus/stubs.rs

nautilus_common/msgbus/
stubs.rs
