use crate::models::{Channel, Symbols};
use serde_json::{json, Value};
#[derive(Debug, Clone)]
pub struct StockSubscription {
pub channel: Channel,
pub symbols: Symbols,
pub intraday_odd_lot: bool,
}
impl StockSubscription {
pub fn new(channel: Channel, symbols: impl Into<Symbols>) -> Self {
Self {
channel,
symbols: symbols.into().normalized(),
intraday_odd_lot: false,
}
}
pub fn with_odd_lot(mut self, odd_lot: bool) -> Self {
self.intraday_odd_lot = odd_lot;
self
}
pub fn keys(&self) -> Vec<String> {
match &self.symbols {
Symbols::Single(s) => vec![self.format_key(s)],
Symbols::Many(v) => v.iter().map(|s| self.format_key(s)).collect(),
}
}
fn format_key(&self, symbol: &str) -> String {
if self.intraday_odd_lot {
format!("{}:{}:oddlot", self.channel.as_str(), symbol)
} else {
format!("{}:{}", self.channel.as_str(), symbol)
}
}
pub fn to_subscribe_data(&self) -> Value {
let mut data = json!({ "channel": self.channel.as_str() });
match &self.symbols {
Symbols::Single(s) => data["symbol"] = json!(s),
Symbols::Many(v) => data["symbols"] = json!(v),
}
if self.intraday_odd_lot {
data["intradayOddLot"] = json!(true);
}
data
}
pub fn to_subscribe_request(&self) -> Value {
json!({
"event": "subscribe",
"data": self.to_subscribe_data()
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn new_single_symbol() {
let sub = StockSubscription::new(Channel::Trades, "2330");
assert_eq!(sub.channel, Channel::Trades);
assert!(matches!(sub.symbols, Symbols::Single(ref s) if s == "2330"));
assert!(!sub.intraday_odd_lot);
}
#[test]
fn new_batch_symbols() {
let sub = StockSubscription::new(Channel::Trades, vec!["2330", "2454"]);
assert!(matches!(sub.symbols, Symbols::Many(ref v) if v == &["2330", "2454"]));
}
#[test]
fn keys_returns_per_symbol_entries() {
let single = StockSubscription::new(Channel::Trades, "2330");
assert_eq!(single.keys(), vec!["trades:2330".to_string()]);
let batch = StockSubscription::new(Channel::Books, ["2330", "2454", "2317"]);
assert_eq!(
batch.keys(),
vec![
"books:2330".to_string(),
"books:2454".to_string(),
"books:2317".to_string()
]
);
}
#[test]
fn odd_lot_modifier_applies_to_all_keys() {
let batch =
StockSubscription::new(Channel::Trades, ["2330", "2454"]).with_odd_lot(true);
assert_eq!(
batch.keys(),
vec![
"trades:2330:oddlot".to_string(),
"trades:2454:oddlot".to_string()
]
);
}
#[test]
fn to_subscribe_data_single_uses_symbol_field() {
let sub = StockSubscription::new(Channel::Candles, "2454");
let data = sub.to_subscribe_data();
assert_eq!(data["channel"], "candles");
assert_eq!(data["symbol"], "2454");
assert!(data.get("symbols").is_none());
assert!(data.get("intradayOddLot").is_none());
}
#[test]
fn to_subscribe_data_batch_uses_symbols_field() {
let sub = StockSubscription::new(Channel::Aggregates, vec!["2330", "0050", "2603"]);
let data = sub.to_subscribe_data();
assert_eq!(data["channel"], "aggregates");
assert_eq!(data["symbols"], json!(["2330", "0050", "2603"]));
assert!(data.get("symbol").is_none());
}
#[test]
fn to_subscribe_data_with_odd_lot() {
let sub = StockSubscription::new(Channel::Trades, "2330").with_odd_lot(true);
let data = sub.to_subscribe_data();
assert_eq!(data["symbol"], "2330");
assert_eq!(data["intradayOddLot"], true);
}
}