#![allow(unused_imports)]
use async_trait::async_trait;
use derive_builder::Builder;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{collections::HashMap, sync::Arc};
use crate::common::{
models::ParamBuildError,
utils::replace_websocket_streams_placeholders,
websocket::{WebsocketBase, WebsocketStream, WebsocketStreams, create_stream_handler},
};
use crate::models::StreamId;
use crate::stocks::websocket_streams::models;
#[async_trait]
pub trait UserStreamsApi: Send + Sync {
async fn order_report_stream(
&self,
params: OrderReportStreamParams,
) -> anyhow::Result<Arc<WebsocketStream<models::OrderReportStreamResponse>>>;
}
pub struct UserStreamsApiClient {
websocket_streams_base: Arc<WebsocketStreams>,
}
impl UserStreamsApiClient {
pub fn new(websocket_streams_base: Arc<WebsocketStreams>) -> Self {
Self {
websocket_streams_base,
}
}
}
#[derive(Clone, Debug, Builder, Deserialize)]
#[builder(pattern = "owned", build_fn(error = "ParamBuildError"))]
pub struct OrderReportStreamParams {
#[builder(setter(into))]
#[serde(rename = "listenKey")]
pub listen_key: String,
#[builder(setter(into), default)]
#[serde(rename = "id", default)]
pub id: Option<String>,
}
impl OrderReportStreamParams {
#[must_use]
pub fn builder(listen_key: String) -> OrderReportStreamParamsBuilder {
OrderReportStreamParamsBuilder::default().listen_key(listen_key)
}
}
#[async_trait]
impl UserStreamsApi for UserStreamsApiClient {
async fn order_report_stream(
&self,
params: OrderReportStreamParams,
) -> anyhow::Result<Arc<WebsocketStream<models::OrderReportStreamResponse>>> {
let OrderReportStreamParams { listen_key, id } = params;
let pairs: &[(&str, Option<String>)] =
&[("listenKey", Some(listen_key.clone())), ("id", id.clone())];
let vars: HashMap<_, _> = pairs
.iter()
.filter_map(|&(k, ref v)| v.clone().map(|v| (k, v)))
.collect();
let id_opt: Option<String> = vars.get("id").map(std::string::ToString::to_string);
let stream = replace_websocket_streams_placeholders("/<listenKey>@orderReport", &vars);
Ok(create_stream_handler::<models::OrderReportStreamResponse>(
WebsocketBase::WebsocketStreams(Arc::clone(&self.websocket_streams_base)),
stream,
id_opt.map(|s| {
if !s.is_empty() && s.bytes().all(|b| b.is_ascii_digit()) {
if let Ok(n) = s.parse::<u32>() {
return StreamId::Number(n);
}
}
StreamId::Str(s)
}),
None,
)
.await)
}
}
#[cfg(all(test, feature = "stocks"))]
mod tests {
use super::*;
use crate::TOKIO_SHARED_RT;
use crate::{
common::websocket::{WebsocketConnection, WebsocketHandler},
config::ConfigurationWebsocketStreams,
};
use serde_json::json;
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::task::yield_now;
async fn make_streams_base() -> (Arc<WebsocketStreams>, Arc<WebsocketConnection>) {
let conn = WebsocketConnection::new("test");
let config = ConfigurationWebsocketStreams::builder()
.build()
.expect("Failed to build configuration");
let streams_base = WebsocketStreams::new(config, vec![conn.clone()], vec![]);
conn.set_handler(streams_base.clone() as Arc<dyn WebsocketHandler>)
.await;
(streams_base, conn)
}
#[test]
fn order_report_stream_should_execute_successfully() {
TOKIO_SHARED_RT.block_on(async {
let (streams_base, _) = make_streams_base().await;
let api = UserStreamsApiClient::new(streams_base.clone());
let id = "test-id-123".to_string();
let params = OrderReportStreamParams::builder(
"pqia91ma19a5s61cv6a81va65sdf19v8a65a1a5s6af0dkfj2a97b8a91d".to_string(),
)
.id(Some(id.clone()))
.build()
.unwrap();
let OrderReportStreamParams { listen_key, id } = params.clone();
let pairs: &[(&str, Option<String>)] =
&[("listenKey", Some(listen_key.clone())), ("id", id.clone())];
let vars: HashMap<_, _> = pairs
.iter()
.filter_map(|&(k, ref v)| v.clone().map(|v| (k, v)))
.collect();
let stream = replace_websocket_streams_placeholders("/<listenKey>@orderReport", &vars);
let ws_stream = api
.order_report_stream(params)
.await
.expect("order_report_stream should return a WebsocketStream");
assert!(
streams_base.is_subscribed(&stream).await,
"expected stream '{stream}' to be subscribed"
);
assert_eq!(ws_stream.id, Some(StreamId::Str("test-id-123".to_string())));
});
}
#[test]
fn order_report_stream_should_handle_incoming_message() {
TOKIO_SHARED_RT.block_on(async {
let (streams_base, conn) = make_streams_base().await;
let api = UserStreamsApiClient::new(streams_base.clone());
let id = "test-id-123".to_string();
let params = OrderReportStreamParams::builder("pqia91ma19a5s61cv6a81va65sdf19v8a65a1a5s6af0dkfj2a97b8a91d".to_string(),).id(Some(id.clone())).build().unwrap();
let OrderReportStreamParams {
listen_key,id,
} = params.clone();
let pairs: &[(&str, Option<String>)] = &[
("listenKey",
Some(listen_key.clone())
),
("id",
id.clone()
),
];
let vars: HashMap<_, _> = pairs
.iter()
.filter_map(|&(k, ref v)| v.clone().map(|v| (k, v)))
.collect();
let stream = replace_websocket_streams_placeholders("/<listenKey>@orderReport", &vars);
let ws_stream = api.order_report_stream(params).await.unwrap();
let called = Arc::new(AtomicBool::new(false));
let called_with_message = called.clone();
ws_stream.on_message(move |_payload: models::OrderReportStreamResponse| {
called_with_message.store(true, Ordering::SeqCst);
});
let payload: Value = serde_json::from_str(r#"{"e":"orderReport","E":1710000000000,"x":"ORDER_UPDATE","i":"b0b6dd9d-8b9b-48a9-ba46-b9d54906e415","ai":"4ce9353c-66d1-46c2-898f-fce867ab0247","b":"EQ_AAPL","q":"USD","S":"buy","o":"limit","p":185.5,"Q":10,"N":1,"fq":5,"FN":927.5,"tc":1856.2,"Z":50,"n":"Regular","s":"partially_filled","T":1710000000000,"U":1710000060000}"#).unwrap_or_else(|_| serde_json::json!({}));
let msg = json!({
"stream": stream,
"data": payload,
});
streams_base.on_message(msg.to_string(), conn.clone()).await;
yield_now().await;
assert!(called.load(Ordering::SeqCst), "expected our callback to have been invoked");
});
}
#[test]
fn order_report_stream_should_not_fire_after_unsubscribe() {
TOKIO_SHARED_RT.block_on(async {
let (streams_base, conn) = make_streams_base().await;
let api = UserStreamsApiClient::new(streams_base.clone());
let id = "test-id-123".to_string();
let params = OrderReportStreamParams::builder("pqia91ma19a5s61cv6a81va65sdf19v8a65a1a5s6af0dkfj2a97b8a91d".to_string(),).id(Some(id.clone())).build().unwrap();
let OrderReportStreamParams {
listen_key,id,
} = params.clone();
let pairs: &[(&str, Option<String>)] = &[
("listenKey",
Some(listen_key.clone())
),
("id",
id.clone()
),
];
let vars: HashMap<_, _> = pairs
.iter()
.filter_map(|&(k, ref v)| v.clone().map(|v| (k, v)))
.collect();
let stream = replace_websocket_streams_placeholders("/<listenKey>@orderReport", &vars);
let ws_stream = api.order_report_stream(params).await.unwrap();
let called = Arc::new(AtomicBool::new(false));
let called_clone = called.clone();
ws_stream.on_message(move |_payload: models::OrderReportStreamResponse| {
called_clone.store(true, Ordering::SeqCst);
});
assert!(streams_base.is_subscribed(&stream).await, "should be subscribed before unsubscribe");
ws_stream.unsubscribe().await;
let payload: Value = serde_json::from_str(r#"{"e":"orderReport","E":1710000000000,"x":"ORDER_UPDATE","i":"b0b6dd9d-8b9b-48a9-ba46-b9d54906e415","ai":"4ce9353c-66d1-46c2-898f-fce867ab0247","b":"EQ_AAPL","q":"USD","S":"buy","o":"limit","p":185.5,"Q":10,"N":1,"fq":5,"FN":927.5,"tc":1856.2,"Z":50,"n":"Regular","s":"partially_filled","T":1710000000000,"U":1710000060000}"#).unwrap_or_else(|_| serde_json::json!({}));
let msg = json!({
"stream": stream,
"data": payload,
});
streams_base.on_message(msg.to_string(), conn.clone()).await;
yield_now().await;
assert!(!called.load(Ordering::SeqCst), "callback should not be invoked after unsubscribe");
});
}
}