use std::collections::HashMap;
use ahash::RandomState;
use async_trait::async_trait;
use indexmap::IndexMap;
use nautilus_model::{
identifiers::InstrumentId,
instruments::{Instrument, InstrumentAny},
};
#[derive(Debug, Default)]
pub struct InstrumentStore {
instruments: IndexMap<InstrumentId, InstrumentAny, RandomState>,
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) -> &IndexMap<InstrumentId, InstrumentAny, RandomState> {
&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_iterates_in_insertion_order() {
let mut store = InstrumentStore::new();
let base = crypto_perpetual_ethusdt();
let ids = [
"SOLUSDT-PERP.BINANCE",
"ADAUSDT-PERP.BINANCE",
"XRPUSDT-PERP.BINANCE",
"BTCUSDT-PERP.BINANCE",
"DOTUSDT-PERP.BINANCE",
"AVAXUSDT-PERP.BINANCE",
];
let expected: Vec<InstrumentId> = ids.iter().map(|id| InstrumentId::from(*id)).collect();
for id in ids {
let mut variant = base.clone();
variant.id = InstrumentId::from(id);
store.add(InstrumentAny::CryptoPerpetual(variant));
}
let keys: Vec<InstrumentId> = store.get_all().keys().copied().collect();
let listed: Vec<InstrumentId> = store.list_all().into_iter().map(Instrument::id).collect();
assert_eq!(keys, expected);
assert_eq!(listed, expected);
}
#[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));
}
}