From 3e062aa43a7f1927284e7ab1243492b320e960c0 Mon Sep 17 00:00:00 2001 From: Godmode Galactus Date: Tue, 14 May 2024 10:21:44 +0200 Subject: [PATCH] Adding tests to test the client --- client/src/client.rs | 87 +++++++++++++++++++++++++++++++ common/src/filters.rs | 4 +- common/src/quic/quinn_reciever.rs | 26 ++------- common/src/types/account.rs | 17 +++++- 4 files changed, 109 insertions(+), 25 deletions(-) diff --git a/client/src/client.rs b/client/src/client.rs index be774b5..0806f3e 100644 --- a/client/src/client.rs +++ b/client/src/client.rs @@ -79,3 +79,90 @@ impl Client { } } } + +#[cfg(test)] +mod tests { + use std::{net::{IpAddr, Ipv4Addr, UdpSocket}, sync::Arc}; + + use futures::StreamExt; + use quic_geyser_common::{filters::{AccountFilter, Filter}, message::Message, quic::{configure_server::configure_server, connection_manager::ConnectionManager}, types::account::Account}; + use quinn::{Endpoint, EndpointConfig, TokioRuntime}; + use solana_sdk::{pubkey::Pubkey, signature::Keypair}; + use tokio::{pin, sync::Notify}; + + use crate::client::Client; + + #[tokio::test] + pub async fn test_client() { + let (config, _) = configure_server( + &Keypair::new(), + IpAddr::V4(Ipv4Addr::LOCALHOST), + 1, + 100000, + 1, + ) + .unwrap(); + + let sock = UdpSocket::bind("0.0.0.0:0").unwrap(); + let port = sock.local_addr().unwrap().port(); + let url = format!("127.0.0.1:{}", port); + let notify_server_start = Arc::new(Notify::new()); + let notify_subscription = Arc::new(Notify::new()); + + let msg_acc_1 = Message::AccountMsg(Account::get_account_for_test(0, 2)); + let msg_acc_2 = Message::AccountMsg(Account::get_account_for_test(1, 20)); + let msg_acc_3 = Message::AccountMsg(Account::get_account_for_test(2, 100)); + let msg_acc_4 = Message::AccountMsg(Account::get_account_for_test(3, 1000)); + let msg_acc_5 = Message::AccountMsg(Account::get_account_for_test(4, 10000)); + let msgs = [msg_acc_1, msg_acc_2, msg_acc_3, msg_acc_4, msg_acc_5]; + + { + let msgs = msgs.clone(); + let notify_server_start = notify_server_start.clone(); + let notify_subscription = notify_subscription.clone(); + tokio::spawn(async move { + let endpoint = Endpoint::new( + EndpointConfig::default(), + Some(config), + sock, + Arc::new(TokioRuntime), + ) + .unwrap(); + + let (connection_manager, _jh) = ConnectionManager::new(endpoint); + notify_server_start.notify_one(); + notify_subscription.notified().await; + for msg in msgs { + connection_manager.dispach(msg, 10).await; + } + }); + } + + notify_server_start.notified().await; + // server started + + let client = Client::new(url, &Keypair::new(), 1024).await.unwrap(); + client.subscribe(vec![Filter::Account(AccountFilter{ + owner: Some(Pubkey::default()), + accounts: None, + })]).await.unwrap(); + + notify_subscription.notify_one(); + + let stream = client.get_stream(); + pin!(stream); + for _ in 0..5 { + let msg = stream.next().await.unwrap(); + match &msg { + Message::AccountMsg(account) => { + let index = account.slot_identifier.slot as usize; + let sent_message = &msgs[index]; + assert_eq!(*sent_message, msg); + }, + _ => { + panic!("should only get account messages") + }, + } + } + } +} diff --git a/common/src/filters.rs b/common/src/filters.rs index ecad576..6ea15e6 100644 --- a/common/src/filters.rs +++ b/common/src/filters.rs @@ -21,8 +21,8 @@ impl Filter { // setting owner to 11111111111111111111111111111111 will subscribe to all the accounts #[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq)] pub struct AccountFilter { - owner: Option, - accounts: Option>, + pub owner: Option, + pub accounts: Option>, } impl AccountFilter { diff --git a/common/src/quic/quinn_reciever.rs b/common/src/quic/quinn_reciever.rs index fe81bed..94e0f38 100644 --- a/common/src/quic/quinn_reciever.rs +++ b/common/src/quic/quinn_reciever.rs @@ -38,7 +38,7 @@ mod tests { }; use quinn::{Endpoint, EndpointConfig, TokioRuntime, VarInt}; - use solana_sdk::{hash::Hash, pubkey::Pubkey, signature::Keypair}; + use solana_sdk::signature::Keypair; use crate::{ message::Message, @@ -46,7 +46,7 @@ mod tests { configure_client::configure_client, configure_server::configure_server, quinn_reciever::recv_message, quinn_sender::send_message, }, - types::{account::Account, slot_identifier::SlotIdentifier}, + types::account::Account, }; #[tokio::test] @@ -71,16 +71,7 @@ mod tests { ) .unwrap(); - let account = Account { - slot_identifier: SlotIdentifier { - slot: 12345678, - blockhash: Hash::new_unique(), - }, - pubkey: Pubkey::new_unique(), - owner: Pubkey::new_unique(), - write_version: 0, - data: vec![6; 2], - }; + let account = Account::get_account_for_test(123456, 2); let message = Message::AccountMsg(account); let jh = { @@ -133,16 +124,7 @@ mod tests { ) .unwrap(); - let account = Account { - slot_identifier: SlotIdentifier { - slot: 12345678, - blockhash: Hash::new_unique(), - }, - pubkey: Pubkey::new_unique(), - owner: Pubkey::new_unique(), - write_version: 0, - data: vec![9; 100_000_000], - }; + let account = Account::get_account_for_test(123456, 100_000_00); let message = Message::AccountMsg(account); let jh = { diff --git a/common/src/types/account.rs b/common/src/types/account.rs index 4ba33c1..946f505 100644 --- a/common/src/types/account.rs +++ b/common/src/types/account.rs @@ -1,5 +1,5 @@ use serde::{Deserialize, Serialize}; -use solana_sdk::pubkey::Pubkey; +use solana_sdk::{clock::Slot, hash::Hash, pubkey::Pubkey}; use super::slot_identifier::SlotIdentifier; @@ -11,3 +11,18 @@ pub struct Account { pub write_version: u64, pub data: Vec, } + +impl Account { + pub fn get_account_for_test(slot: Slot, data_size: usize) -> Self { + Account { + slot_identifier: SlotIdentifier { + slot, + blockhash: Hash::new_unique(), + }, + pubkey: Pubkey::new_unique(), + owner: Pubkey::new_unique(), + write_version: 0, + data: vec![178; data_size], + } + } +} \ No newline at end of file