use std::collections::HashMap;
use ahash::AHashMap;
use async_trait::async_trait;
use nautilus_model::{
identifiers::InstrumentId,
instruments::{Instrument, InstrumentAny},
};
#[derive(Debug, Default)]
pub struct InstrumentStore {
instruments: AHashMap<InstrumentId, InstrumentAny>,
initialized: bool,
}
impl InstrumentStore {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn add(&mut self, instrument: InstrumentAny) {
self.instruments.insert(instrument.id(), instrument);
}
pub fn add_bulk(&mut self, instruments: Vec<InstrumentAny>) {
for instrument in instruments {
self.add(instrument);
}
}
#[must_use]
pub fn find(&self, instrument_id: &InstrumentId) -> Option<&InstrumentAny> {
self.instruments.get(instrument_id)
}
#[must_use]
pub fn contains(&self, instrument_id: &InstrumentId) -> bool {
self.instruments.contains_key(instrument_id)
}
#[must_use]
pub fn get_all(&self) -> &AHashMap<InstrumentId, InstrumentAny> {
&self.instruments
}
#[must_use]
pub fn list_all(&self) -> Vec<&InstrumentAny> {
self.instruments.values().collect()
}
#[must_use]
pub fn count(&self) -> usize {
self.instruments.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.instruments.is_empty()
}
#[must_use]
pub fn is_initialized(&self) -> bool {
self.initialized
}
pub fn set_initialized(&mut self) {
self.initialized = true;
}
pub fn clear(&mut self) {
self.instruments.clear();
self.initialized = false;
}
}
#[async_trait(?Send)]
pub trait InstrumentProvider {
fn store(&self) -> &InstrumentStore;
fn store_mut(&mut self) -> &mut InstrumentStore;
async fn load_all(&mut self, filters: Option<&HashMap<String, String>>) -> anyhow::Result<()>;
async fn load_ids(
&mut self,
instrument_ids: &[InstrumentId],
filters: Option<&HashMap<String, String>>,
) -> anyhow::Result<()> {
for instrument_id in instrument_ids {
self.load(instrument_id, filters).await?;
}
Ok(())
}
async fn load(
&mut self,
instrument_id: &InstrumentId,
filters: Option<&HashMap<String, String>>,
) -> anyhow::Result<()>;
}
#[cfg(test)]
mod tests {
use nautilus_model::instruments::{InstrumentAny, stubs::crypto_perpetual_ethusdt};
use rstest::rstest;
use super::*;
#[rstest]
fn test_instrument_store_default_is_empty() {
let store = InstrumentStore::new();
assert!(store.is_empty());
assert_eq!(store.count(), 0);
assert!(!store.is_initialized());
}
#[rstest]
fn test_instrument_store_add_and_find() {
let mut store = InstrumentStore::new();
let instrument = InstrumentAny::CryptoPerpetual(crypto_perpetual_ethusdt());
let id = instrument.id();
store.add(instrument);
assert_eq!(store.count(), 1);
assert!(!store.is_empty());
assert!(store.contains(&id));
assert!(store.find(&id).is_some());
}
#[rstest]
fn test_instrument_store_add_bulk() {
let mut store = InstrumentStore::new();
let instrument = InstrumentAny::CryptoPerpetual(crypto_perpetual_ethusdt());
let id = instrument.id();
store.add_bulk(vec![instrument]);
assert_eq!(store.count(), 1);
assert!(store.contains(&id));
}
#[rstest]
fn test_instrument_store_get_all() {
let mut store = InstrumentStore::new();
let instrument = InstrumentAny::CryptoPerpetual(crypto_perpetual_ethusdt());
store.add(instrument);
let all = store.get_all();
assert_eq!(all.len(), 1);
}
#[rstest]
fn test_instrument_store_list_all() {
let mut store = InstrumentStore::new();
let instrument = InstrumentAny::CryptoPerpetual(crypto_perpetual_ethusdt());
store.add(instrument);
let list = store.list_all();
assert_eq!(list.len(), 1);
}
#[rstest]
fn test_instrument_store_clear() {
let mut store = InstrumentStore::new();
let instrument = InstrumentAny::CryptoPerpetual(crypto_perpetual_ethusdt());
store.add(instrument);
store.set_initialized();
assert!(store.is_initialized());
assert_eq!(store.count(), 1);
store.clear();
assert!(!store.is_initialized());
assert!(store.is_empty());
}
#[rstest]
fn test_instrument_store_find_missing_returns_none() {
let store = InstrumentStore::new();
let id = InstrumentId::from("UNKNOWN-UNKNOWN.VENUE");
assert!(store.find(&id).is_none());
assert!(!store.contains(&id));
}
#[rstest]
fn test_instrument_store_add_replaces_existing() {
let mut store = InstrumentStore::new();
let instrument1 = InstrumentAny::CryptoPerpetual(crypto_perpetual_ethusdt());
let instrument2 = InstrumentAny::CryptoPerpetual(crypto_perpetual_ethusdt());
let id = instrument1.id();
store.add(instrument1);
store.add(instrument2);
assert_eq!(store.count(), 1);
assert!(store.contains(&id));
}
}