diff --git a/tests/ws_integration_tests.rs b/tests/ws_integration_tests.rs index a609eb3..aa3c405 100644 --- a/tests/ws_integration_tests.rs +++ b/tests/ws_integration_tests.rs @@ -9,10 +9,25 @@ use futures::StreamExt; use polyfill_rs::{ClobClient, OrderBookManager, WebSocketStream, WsBookUpdateProcessor}; +use std::env; use std::time::Duration; const HOST: &str = "https://clob.polymarket.com"; const WS_MARKET_URL: &str = "wss://ws-subscriptions-clob.polymarket.com/ws/market"; +const WS_USER_URL: &str = "wss://ws-subscriptions-clob.polymarket.com/ws/user"; +const CHAIN_ID: u64 = 137; + +fn load_private_key() -> String { + dotenvy::dotenv().ok(); + env::var("POLYMARKET_PRIVATE_KEY").expect("POLYMARKET_PRIVATE_KEY must be set (env or .env)") +} + +fn stability_secs(default_secs: u64) -> u64 { + env::var("POLYFILL_WS_STABILITY_SECS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(default_secs) +} #[tokio::test(flavor = "multi_thread")] #[ignore] @@ -62,3 +77,107 @@ async fn test_real_ws_market_book_applier_receives_book_update() { "expected some book levels after applying an update" ); } + +#[tokio::test(flavor = "multi_thread")] +#[ignore] +async fn test_real_ws_market_book_applier_connection_stable() { + // Subscribe to a handful of active tokens and keep the connection open for a while. + // We don't require constant message flow; we just fail on close or error. + let client = ClobClient::new(HOST); + let markets = client + .get_sampling_markets(None) + .await + .expect("failed to fetch markets"); + + let token_ids: Vec = markets + .data + .iter() + .filter(|m| m.active && !m.closed) + .filter_map(|m| m.tokens.first().map(|t| t.token_id.clone())) + .take(10) + .collect(); + + assert!(!token_ids.is_empty(), "no active token IDs found"); + + let books = OrderBookManager::new(256); + for token_id in &token_ids { + books + .get_or_create_book(token_id) + .expect("failed to create book"); + } + + let mut ws = WebSocketStream::new(WS_MARKET_URL); + ws.subscribe_market_channel(token_ids.clone()) + .await + .expect("failed to subscribe market channel"); + + let processor = WsBookUpdateProcessor::new(256 * 1024); + let mut applier = ws.into_book_applier(&books, processor); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(stability_secs(15)); + let mut saw_book_message = false; + + while tokio::time::Instant::now() < deadline { + match tokio::time::timeout(Duration::from_secs(1), applier.next()).await { + Ok(Some(Ok(stats))) => { + if stats.book_messages > 0 { + saw_book_message = true; + } + }, + Ok(Some(Err(e))) => panic!("WS processing error: {:?}", e), + Ok(None) => panic!("WS stream ended unexpectedly"), + Err(_) => { + // No message in this interval; keep waiting. + }, + } + } + + assert!( + saw_book_message, + "did not observe any `book` messages during stability window" + ); +} + +#[tokio::test(flavor = "multi_thread")] +#[ignore] +async fn test_real_ws_user_channel_connection_stable() { + // Connect + subscribe to the user channel and keep the connection open for a while. + // + // Note: depending on account activity, the WS may not emit any `order`/`trade` messages + // during the window. This test primarily asserts we can authenticate + subscribe + // and that the connection doesn't immediately drop. + let private_key = load_private_key(); + let auth_client = ClobClient::with_l1_headers(HOST, &private_key, CHAIN_ID); + let api_creds = auth_client + .create_or_derive_api_key(None) + .await + .expect("failed to create/derive api key"); + + let markets = auth_client + .get_sampling_markets(None) + .await + .expect("failed to fetch markets"); + let market_id = markets + .data + .iter() + .find(|m| m.active && !m.closed) + .map(|m| m.condition_id.clone()) + .expect("no active markets found"); + + let mut ws = WebSocketStream::new(WS_USER_URL).with_auth(api_creds); + ws.subscribe_user_channel(vec![market_id]) + .await + .expect("failed to subscribe user channel"); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(stability_secs(10)); + while tokio::time::Instant::now() < deadline { + match tokio::time::timeout(Duration::from_secs(1), ws.next()).await { + Ok(Some(Ok(_msg))) => {}, + Ok(Some(Err(e))) => panic!("WS error: {:?}", e), + Ok(None) => panic!("WS stream ended unexpectedly"), + Err(_) => { + // No message in this interval; keep waiting. + }, + } + } +}