vyper_client_rs/
websocket.rs

1use std::sync::Arc;
2use tokio::sync::Mutex;
3use async_trait::async_trait;
4use serde::{Serialize, Deserialize};
5use serde_json::Value;
6use tokio_tungstenite::{connect_async, tungstenite::protocol::Message, WebSocketStream, MaybeTlsStream};
7use futures_util::{SinkExt, StreamExt};
8use tokio::net::TcpStream;
9use crate::types::{ChainAction, TokenPair};
10use crate::errors::VyperError;  
11
12#[derive(Debug, Clone, PartialEq)]
13pub enum FeedType {
14    TokenEvents,
15    MigrationEvents,
16    WalletEvents,
17}
18
19impl ToString for FeedType {
20    fn to_string(&self) -> String {
21        match self {
22            FeedType::TokenEvents => "token-events".to_string(),
23            FeedType::MigrationEvents => "migration-events".to_string(),
24            FeedType::WalletEvents => "wallet-events".to_string(),
25        }
26    }
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize)]
30#[serde(rename_all = "lowercase")]
31pub enum SubscriptionMessageType {
32    Subscribe,
33    Unsubscribe,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize)]
37#[serde(rename_all = "PascalCase")]
38pub enum SubscriptionType {
39    PumpfunTokens,
40    RaydiumAmmTokens,
41    RaydiumCpmmTokens,
42    RaydiumClmmTokens,
43}
44
45#[derive(Debug, Clone, Serialize)]
46pub struct TokenSubscriptionMessage {
47    pub action: SubscriptionMessageType,
48    pub types: Vec<SubscriptionType>,
49}
50
51#[derive(Debug, Clone, Serialize)]
52pub struct WalletSubscriptionMessage {
53    pub action: SubscriptionMessageType,
54    pub wallets: Vec<String>,
55}
56
57type MessageHandler = Arc<Mutex<dyn FnMut(Value) + Send + Sync>>;
58
59#[async_trait]
60pub trait WebSocketConnection: Send + Sync {
61    async fn connect(&mut self, url: &str) -> Result<(), VyperError>;
62    async fn send(&mut self, data: &str) -> Result<(), VyperError>;
63    async fn receive(&mut self) -> Result<String, VyperError>;
64    async fn close(&mut self) -> Result<(), VyperError>;
65}
66
67pub struct VyperWebsocketClient {
68    base_url: String,
69    api_key: String,
70    conn: Arc<Mutex<Option<Box<dyn WebSocketConnection>>>>,
71    message_handler: Option<MessageHandler>,
72    current_feed_type: Arc<Mutex<Option<FeedType>>>,
73}
74
75impl VyperWebsocketClient {
76    pub fn new(api_key: String) -> Self {
77        Self {
78            base_url: "wss://api.vyper.trade/api/v1/ws".to_string(),
79            api_key,
80            conn: Arc::new(Mutex::new(None)),
81            message_handler: None,
82            current_feed_type: Arc::new(Mutex::new(None)),
83        }
84    }
85
86    pub async fn connect(&self, feed_type: FeedType) -> Result<(), VyperError> {
87        let mut conn_guard = self.conn.lock().await;
88        let url = format!("{}/{}?apiKey={}", self.base_url, feed_type.to_string(), self.api_key);
89
90        if conn_guard.is_none() {
91            let mut new_conn: Box<dyn WebSocketConnection> = Box::new(WebSocketImpl::new());
92            new_conn.connect(&url).await.map_err(|e| VyperError::WebsocketError {
93                message: e.to_string(),
94                status_code: None,
95                connection_info: Some(url.clone()),
96            })?;
97            *conn_guard = Some(new_conn);
98        } else {
99            conn_guard.as_mut().unwrap().connect(&url).await.map_err(|e| VyperError::WebsocketError {
100                message: e.to_string(),
101                status_code: None,
102                connection_info: Some(url.clone()),
103            })?;
104        }
105
106        let mut feed_type_guard = self.current_feed_type.lock().await;
107        *feed_type_guard = Some(feed_type);
108
109        Ok(())
110    }
111
112    pub async fn subscribe<M: Serialize>(&self, feed_type: FeedType, message: M) -> Result<(), VyperError> {
113        let mut conn_guard = self.conn.lock().await;
114        let feed_type_guard = self.current_feed_type.lock().await;
115
116        if conn_guard.is_none() {
117            return Err(VyperError::WebsocketError {
118                message: "Not connected".to_string(),
119                status_code: None,
120                connection_info: None,
121            });
122        }
123
124        if feed_type != *feed_type_guard.as_ref().unwrap() {
125            return Err(VyperError::WebsocketError {
126                message: "Feed type mismatch".to_string(),
127                status_code: None,
128                connection_info: None,
129            });
130        }
131
132        let data = serde_json::to_string(&message).map_err(|e| VyperError::DeserializeError(e))?;
133        conn_guard.as_mut().unwrap().send(&data).await.map_err(|e| VyperError::WebsocketError {
134            message: e.to_string(),
135            status_code: None,
136            connection_info: None,
137        })?;
138
139        Ok(())
140    }
141
142    pub async fn unsubscribe<M: Serialize>(&self, feed_type: FeedType, message: M) -> Result<(), VyperError> {
143        self.subscribe(feed_type, message).await
144    }
145
146    pub async fn listen(&self) -> Result<(), VyperError> {
147        let mut conn_guard = self.conn.lock().await;
148        let feed_type_guard = self.current_feed_type.lock().await;
149
150        if conn_guard.is_none() {
151            return Err(VyperError::WebsocketError {
152                message: "Not connected".to_string(),
153                status_code: None,
154                connection_info: None,
155            });
156        }
157
158        while let Ok(msg) = conn_guard.as_mut().unwrap().receive().await {
159            let raw_data: Value = serde_json::from_str(&msg).map_err(|e| VyperError::DeserializeError(e))?;
160            let converted_data = self.convert_message(&raw_data, &feed_type_guard).await;
161
162            if let Err(_) = converted_data {
163                continue;
164            }
165
166            if let Some(ref handler) = self.message_handler {
167                let mut handler = handler.lock().await;
168                handler(converted_data.unwrap());
169            }
170        }
171
172        Ok(())
173    }
174
175    async fn convert_message(&self, data: &Value, feed_type: &Option<FeedType>) -> Result<Value, VyperError> {
176        if feed_type.is_none() {
177            return Err(VyperError::WebsocketError {
178                message: "Feed type is not set".to_string(),
179                status_code: None,
180                connection_info: None,
181            });
182        }
183    
184        match feed_type.as_ref().unwrap() {
185            FeedType::WalletEvents => {
186                let chain_action: ChainAction = serde_json::from_value(data.clone()).map_err(|e| VyperError::DeserializeError(e))?;
187                Ok(serde_json::to_value(chain_action).map_err(|e| VyperError::DeserializeError(e))?)
188            },
189            FeedType::MigrationEvents | FeedType::TokenEvents => {
190                let token_pair: TokenPair = serde_json::from_value(data.clone()).map_err(|e| VyperError::DeserializeError(e))?;
191                Ok(serde_json::to_value(token_pair).map_err(|e| VyperError::DeserializeError(e))?)
192            },
193        }
194    }
195
196    pub async fn disconnect(&self) -> Result<(), VyperError> {
197        let mut conn_guard = self.conn.lock().await;
198        if let Some(ref mut ws) = conn_guard.as_mut() {
199            ws.close().await.map_err(|e| VyperError::WebsocketError {
200                message: e.to_string(),
201                status_code: None,
202                connection_info: None,
203            })?;
204        }
205
206        *conn_guard = None;
207        let mut feed_type_guard = self.current_feed_type.lock().await;
208        *feed_type_guard = None;
209
210        Ok(())
211    }
212
213    pub async fn ping(&self) -> Result<(), VyperError> {
214        let mut conn_guard = self.conn.lock().await;
215        if let Some(ref mut ws) = conn_guard.as_mut() {
216            ws.send("ping").await.map_err(|e| VyperError::WebsocketError {
217                message: e.to_string(),
218                status_code: None,
219                connection_info: None,
220            })?;
221        }
222        Ok(())
223    }
224
225    pub fn set_message_handler<F>(&mut self, handler: F)
226    where
227        F: Fn(Value) + Send + Sync + 'static,
228    {
229        self.message_handler = Some(Arc::new(Mutex::new(handler)));
230    }
231}
232
233struct WebSocketImpl {
234    ws_stream: Option<WebSocketStream<MaybeTlsStream<TcpStream>>>,
235}
236
237impl WebSocketImpl {
238    fn new() -> Self {
239        Self {
240            ws_stream: None,
241        }
242    }
243}
244
245#[async_trait]
246impl WebSocketConnection for WebSocketImpl {
247    async fn connect(&mut self, url: &str) -> Result<(), VyperError> {
248        let (ws_stream, _) = connect_async(url).await.map_err(|e| VyperError::WebsocketError {
249            message: e.to_string(),
250            status_code: None,
251            connection_info: Some(url.to_string()),
252        })?;
253        self.ws_stream = Some(ws_stream);
254        Ok(())
255    }
256
257    async fn send(&mut self, data: &str) -> Result<(), VyperError> {
258        if let Some(ws_stream) = &mut self.ws_stream {
259            ws_stream.send(Message::Text(data.to_string())).await.map_err(|e| VyperError::WebsocketError {
260                message: e.to_string(),
261                status_code: None,
262                connection_info: None,
263            })?;
264            Ok(())
265        } else {
266            Err(VyperError::WebsocketError {
267                message: "WebSocket is not connected".to_string(),
268                status_code: None,
269                connection_info: None,
270            })
271        }
272    }
273
274    async fn receive(&mut self) -> Result<String, VyperError> {
275        if let Some(ws_stream) = &mut self.ws_stream {
276            if let Some(message) = ws_stream.next().await {
277                match message {
278                    Ok(message) => match message {
279                        Message::Text(text) => Ok(text),
280                        _ => Err(VyperError::WebsocketError {
281                            message: "Received non-text message".to_string(),
282                            status_code: None,
283                            connection_info: None,
284                        }),
285                    },
286                    Err(e) => Err(VyperError::WebsocketError {
287                        message: e.to_string(),
288                        status_code: None,
289                        connection_info: None,
290                    }),
291                }
292            } else {
293                Err(VyperError::WebsocketError {
294                    message: "WebSocket stream ended".to_string(),
295                    status_code: None,
296                    connection_info: None,
297                })
298            }
299        } else {
300            Err(VyperError::WebsocketError {
301                message: "WebSocket is not connected".to_string(),
302                status_code: None,
303                connection_info: None,
304            })
305        }
306    }
307
308    async fn close(&mut self) -> Result<(), VyperError> {
309        if let Some(ws_stream) = &mut self.ws_stream {
310            ws_stream.close(None).await.map_err(|e| VyperError::WebsocketError {
311                message: e.to_string(),
312                status_code: None,
313                connection_info: None,
314            })?;
315            self.ws_stream = None;
316            Ok(())
317        } else {
318            Err(VyperError::WebsocketError {
319                message: "WebSocket is not connected".to_string(),
320                status_code: None,
321                connection_info: None,
322            })
323        }
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330    use mockall::predicate::*;
331    use mockall::mock;
332    use std::sync::Arc;
333
334    mock! {
335        WebSocket {}
336        #[async_trait]
337        impl WebSocketConnection for WebSocket {
338            async fn connect(&mut self, url: &str) -> Result<(), VyperError>;
339            async fn send(&mut self, data: &str) -> Result<(), VyperError>;
340            async fn receive(&mut self) -> Result<String, VyperError>;
341            async fn close(&mut self) -> Result<(), VyperError>;
342        }
343    }    
344
345    #[tokio::test]
346    async fn test_new_client() {
347        let client = VyperWebsocketClient::new("test_api_key".to_string());
348        assert_eq!(client.base_url, "wss://api.vyper.trade/api/v1/ws");
349        assert_eq!(client.api_key, "test_api_key");
350    }
351
352    #[tokio::test]
353    async fn test_connect() {
354        let mut mock_ws = MockWebSocket::new();
355        
356        mock_ws.expect_connect()
357            .with(eq("wss://api.vyper.trade/api/v1/ws/token-events?apiKey=test_api_key"))
358            .times(1)
359            .returning(|_| Ok(()));
360
361        let client = VyperWebsocketClient::new("test_api_key".to_string());
362
363        {
364            let mut conn_guard = client.conn.lock().await;
365            *conn_guard = Some(Box::new(mock_ws));
366        }
367
368        let result = client.connect(FeedType::TokenEvents).await;
369        assert!(result.is_ok());
370
371        let feed_type_guard = client.current_feed_type.lock().await;
372        assert_eq!(*feed_type_guard, Some(FeedType::TokenEvents));
373    }
374
375    #[tokio::test]
376    async fn test_disconnect() {
377        let mut mock_ws = MockWebSocket::new();
378        mock_ws.expect_close()
379            .returning(|| Ok(()));
380
381        let client = VyperWebsocketClient::new("test_api_key".to_string());
382        *client.conn.lock().await = Some(Box::new(mock_ws));
383        
384        let result = client.disconnect().await;
385        assert!(result.is_ok());
386    }
387
388    #[tokio::test]
389    async fn test_subscribe() {
390        let mut mock_ws = MockWebSocket::new();
391        mock_ws.expect_send()
392            .returning(|_| Ok(()));
393
394        let client = VyperWebsocketClient::new("test_api_key".to_string());
395        *client.conn.lock().await = Some(Box::new(mock_ws));
396        *client.current_feed_type.lock().await = Some(FeedType::TokenEvents);
397
398        let message = TokenSubscriptionMessage {
399            action: SubscriptionMessageType::Subscribe,
400            types: vec![SubscriptionType::PumpfunTokens],
401        };
402
403        let result = client.subscribe(FeedType::TokenEvents, message).await;
404        assert!(result.is_ok());
405    }
406
407    #[tokio::test]
408    async fn test_unsubscribe() {
409        let mut mock_ws = MockWebSocket::new();
410        mock_ws.expect_send()
411            .returning(|_| Ok(()));
412
413        let client = VyperWebsocketClient::new("test_api_key".to_string());
414        *client.conn.lock().await = Some(Box::new(mock_ws));
415        *client.current_feed_type.lock().await = Some(FeedType::TokenEvents);
416
417        let message = TokenSubscriptionMessage {
418            action: SubscriptionMessageType::Unsubscribe,
419            types: vec![SubscriptionType::PumpfunTokens],
420        };
421
422        let result = client.unsubscribe(FeedType::TokenEvents, message).await;
423        assert!(result.is_ok());
424    }
425
426    #[tokio::test]
427    async fn test_ping() {
428        let mut mock_ws = MockWebSocket::new();
429        mock_ws.expect_send()
430            .with(eq("ping"))
431            .returning(|_| Ok(()));
432
433        let client = VyperWebsocketClient::new("test_api_key".to_string());
434        *client.conn.lock().await = Some(Box::new(mock_ws));
435        
436        let result = client.ping().await;
437        assert!(result.is_ok());
438    }
439
440    #[tokio::test]
441    async fn test_listen() {
442        let mut mock_ws = MockWebSocket::new();
443        mock_ws.expect_receive()
444            .returning(|| Ok(r#"{
445                "event": "tokenPair",
446                "chainId": 1,
447                "tokenMint": "0x123",
448                "name": "Test Token",
449                "symbol": "TEST",
450                "buyTxnCount": 10,
451                "sellTxnCount": 5,
452                "tokenPriceUsd": 1.5,
453                "tokenPriceAsset": 0.1,
454                "volumeUsd": 1000.0,
455                "volumeAsset": 100.0,
456                "tokenMarketCapUsd": 1000000.0,
457                "tokenMarketCapAsset": 100000.0,
458                "tokenLiquidityUsd": 500000.0,
459                "tokenLiquidityAsset": 50000.0,
460                "transactionCount": 15,
461                "contractCreator": "0x456",
462                "lpCreator": "0x789",
463                "createdTimestamp": 1631234567,
464                "totalSupply": 1000000.0,
465                "pooledToken": 500000.0,
466                "pooledAsset": 50000.0,
467                "initialUsdLiquidity": 100000.0,
468                "initialAssetLiquidity": 10000.0,
469                "priceChangePercent": 5.0,
470                "top10HoldingPercent": 60.0,
471                "lpBurned": false,
472                "tokenType": "SPL",
473                "marketId": "market1"
474            }"#.to_string()));
475        mock_ws.expect_receive()
476            .returning(|| Err("Connection closed".into()));
477
478        let mut client = VyperWebsocketClient::new("test_api_key".to_string());
479        *client.conn.lock().await = Some(Box::new(mock_ws));
480        *client.current_feed_type.lock().await = Some(FeedType::TokenEvents);
481        
482        let received_message = Arc::new(Mutex::new(false));
483        let received_message_clone = received_message.clone();
484        
485        client.set_message_handler(move |msg| {
486            let received_message_clone = received_message_clone.clone();
487            tokio::spawn(async move {
488                assert!(msg.get("chainId").is_some(), "chainId field is missing");
489                assert!(msg.get("tokenMint").is_some(), "tokenMint field is missing");
490                assert!(msg.get("name").is_some(), "name field is missing");
491                assert!(msg.get("symbol").is_some(), "symbol field is missing");
492                let mut received = received_message_clone.lock().await;
493                *received = true;
494            });
495        });
496
497        let listen_future = client.listen();
498        tokio::time::timeout(std::time::Duration::from_secs(1), listen_future).await.unwrap_or(Ok(())).unwrap();
499        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
500
501        let received = received_message.lock().await;
502        assert!(*received, "Message was not received and processed");
503    }
504
505    #[tokio::test]
506    async fn test_error_handling() {
507        let mut mock_ws = MockWebSocket::new();
508        mock_ws.expect_connect()
509            .returning(|_| Err("Connection failed".into()));
510
511        let client = VyperWebsocketClient::new("test_api_key".to_string());
512        *client.conn.lock().await = Some(Box::new(mock_ws));
513        
514        let result = client.connect(FeedType::TokenEvents).await;
515        assert!(result.is_err());
516    }
517
518    #[tokio::test]
519    async fn test_set_message_handler() {
520        let client = VyperWebsocketClient::new("test_api_key".to_string());
521        let mut client = client;
522        client.set_message_handler(|_| {
523            // Do nothing, just testing if it sets without error
524        });
525        assert!(client.message_handler.is_some());
526    }
527
528    #[tokio::test]
529    async fn test_feed_type_mismatch() {
530        let mut mock_ws = MockWebSocket::new();
531        mock_ws.expect_send()
532            .returning(|_| Ok(()));
533
534        let client = VyperWebsocketClient::new("test_api_key".to_string());
535        *client.conn.lock().await = Some(Box::new(mock_ws));
536        *client.current_feed_type.lock().await = Some(FeedType::TokenEvents);
537
538        let message = WalletSubscriptionMessage {
539            action: SubscriptionMessageType::Subscribe,
540            wallets: vec!["0x123".to_string()],
541        };
542
543        let result = client.subscribe(FeedType::WalletEvents, message).await;
544        assert!(result.is_err());
545    }
546}
vyper_client_rs/websocket.rs

vyper_client_rs/
websocket.rs
