Skip to content

Commit

Permalink
add more tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@dharanad
dharanad committed Jul 31, 2024
1 parent ab19e62 commit 8530aa7
Show file tree
Hide file tree
Showing 3 changed files with 67 additions and 20 deletions.
1 change: 0 additions & 1 deletion Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ hex = "0.4"
sha2 = "0.10.8"

[dev-dependencies]
tempdir = "0.3.7"
tempfile = "3.10.1"

[[example]]
Expand Down
26 changes: 12 additions & 14 deletions src/network.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,6 @@ impl NetworkLayer for TCPManager {
let listener = TcpListener::bind(addr).await?;
*self.listener.lock().await = Some(listener);
*is_open = true;
println!("Listening on {}", addr);
Ok(())
}

Expand All @@ -126,7 +125,6 @@ impl NetworkLayer for TCPManager {
}
*self.listener.lock().await = None;
*is_open = false;
println!("Listener closed");
Ok(())
}
}
Expand Down Expand Up @@ -156,42 +154,42 @@ mod tests {

#[tokio::test]
async fn test_send_closed_connection() {
let network = TCPManager::new(LOCALHOST.to_string(), 8082);
let network = TCPManager::new(LOCALHOST.to_string(), 8020);
let data = vec![1, 2, 3];
network.open().await.unwrap();
let network_clone = network.clone();
let _ = tokio::spawn(async move {
tokio::spawn(async move {
let _ = network_clone.receive().await.unwrap();
});

let send_result = network.send(LOCALHOST, "8083", &data).await;
let send_result = network.send(LOCALHOST, "8021", &data).await;
assert!(send_result.is_err());
}

#[tokio::test]
async fn test_receive_happy_case() {
let network = TCPManager::new(LOCALHOST.to_string(), 8082);
let network = TCPManager::new(LOCALHOST.to_string(), 8030);
let data = vec![1, 2, 3];
network.open().await.unwrap();
let network_clone = network.clone();
let handler = tokio::spawn(async move { network_clone.receive().await.unwrap() });

network.send(LOCALHOST, "8082", &data).await.unwrap();
network.send(LOCALHOST, "8030", &data).await.unwrap();
let rx_data = handler.await.unwrap();
assert_eq!(rx_data, data)
}

#[tokio::test]
async fn test_open() {
let network = TCPManager::new(LOCALHOST.to_string(), 8081);
let network = TCPManager::new(LOCALHOST.to_string(), 8040);
let status = network.open().await;
assert!(status.is_ok());
assert!(*network.is_open.lock().await);
}

#[tokio::test]
async fn test_reopen_opened_port() {
let network = TCPManager::new(LOCALHOST.to_string(), 8081);
let network = TCPManager::new(LOCALHOST.to_string(), 8042);
let status = network.open().await;
assert!(status.is_ok());
let another_network = network.clone();
Expand All @@ -201,7 +199,7 @@ mod tests {

#[tokio::test]
async fn test_close() {
let network = TCPManager::new(LOCALHOST.to_string(), 8081);
let network = TCPManager::new(LOCALHOST.to_string(), 8046);
let _ = network.open().await;

let close_status = network.close().await;
Expand All @@ -213,7 +211,7 @@ mod tests {
async fn test_broadcast_happy_case() {
let data = vec![1, 2, 3, 4];
// Node which is about to broadcast data
let broadcasting_node = TCPManager::new(LOCALHOST.to_string(), 8081);
let broadcasting_node = TCPManager::new(LOCALHOST.to_string(), 8050);
broadcasting_node.open().await.unwrap();
assert!(*broadcasting_node.is_open.lock().await);

Expand All @@ -222,7 +220,7 @@ mod tests {
// vec to keep track of the address for receiver nodes
let mut receiver_addresses = vec![];

for p in 8082..8090 {
for p in 8051..8060 {
// Create a receiver node
let rx = TCPManager::new(LOCALHOST.to_string(), p);
receiver_addresses.push(format!("{}:{}", LOCALHOST, p));
Expand Down Expand Up @@ -257,15 +255,15 @@ mod tests {
async fn test_broadcast_some_nodes_down() {
let data = vec![1, 2, 3, 4];
// Node which is about to broadcast data
let broadcasting_node = TCPManager::new(LOCALHOST.to_string(), 8081);
let broadcasting_node = TCPManager::new(LOCALHOST.to_string(), 8061);
broadcasting_node.open().await.unwrap();
assert!(*broadcasting_node.is_open.lock().await);

// vec to keep track of all other nodes which should be receiving data
let mut receivers = vec![];
// vec to keep track of the address for receiver nodes
let mut receiver_addresses = vec![];
for p in 8082..8092 {
for p in 8062..8070 {
// Create a receiver node
let rx = TCPManager::new(LOCALHOST.to_string(), p);
receiver_addresses.push(format!("{}:{}", LOCALHOST, p));
Expand Down
60 changes: 55 additions & 5 deletions src/storage.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,10 @@ impl LocalStorage {
let mut buffer = Vec::new();
file.read_to_end(&mut buffer).await?;

if buffer.is_empty() {
return Err("File is empty".into());
}

if buffer.len() < CHECKSUM_LEN {
return Err("File is potentially malicious".into());
}
Expand Down Expand Up @@ -141,12 +145,12 @@ fn retrieve_checksum(data: &[u8]) -> [u8; CHECKSUM_LEN] {

#[cfg(test)]
mod tests {
use std::io::Read;
use std::io::{Read, Seek, SeekFrom, Write};

use tempfile::NamedTempFile;

use crate::storage::{
calculate_checksum, CHECKSUM_LEN, LocalStorage, retrieve_checksum, Storage,
calculate_checksum, retrieve_checksum, LocalStorage, Storage, CHECKSUM_LEN,
};

#[test]
Expand Down Expand Up @@ -190,10 +194,11 @@ mod tests {

#[tokio::test]
async fn test_compaction_file_lt_max_file_size() {
let mut tmp_file = NamedTempFile::new().unwrap();
let mut storage: Box<dyn Storage> = Box::new(LocalStorage::new_from_path(tmp_file.path()));
let tmp_file = NamedTempFile::new().unwrap();
let storage: Box<dyn Storage> = Box::new(LocalStorage::new_from_path(tmp_file.path()));
let mock_data = vec![0u8; 1_000_000 /*1 MB*/ - 500];
let store_result = storage.store(&mock_data).await;
assert!(store_result.is_ok());

tmp_file.as_file().sync_all().unwrap();

Expand All @@ -205,10 +210,11 @@ mod tests {

#[tokio::test]
async fn test_compaction_file_gt_max_file_size() {
let mut tmp_file = NamedTempFile::new().unwrap();
let tmp_file = NamedTempFile::new().unwrap();
let storage: Box<dyn Storage> = Box::new(LocalStorage::new_from_path(tmp_file.path()));
let mock_data = vec![0u8; 1_000_000 /*1 MB*/];
let store_result = storage.store(&mock_data).await;
assert!(store_result.is_ok());

tmp_file.as_file().sync_all().unwrap();

Expand All @@ -217,4 +223,48 @@ mod tests {

assert!(!tmp_file.path().exists());
}

#[tokio::test]
async fn test_retrieve() {
let tmp_file = NamedTempFile::new().unwrap();
let storage: Box<dyn Storage> = Box::new(LocalStorage::new_from_path(tmp_file.path()));
let test_data = "Some mocked data".as_bytes();

storage.store(test_data).await.unwrap();

let retrieved_result = storage.retrieve().await;
assert!(retrieved_result.is_ok());

assert_eq!(test_data, retrieved_result.unwrap());
}

#[tokio::test]
async fn test_turned_malicious_file_corrupted() {
let mut tmp_file = NamedTempFile::new().unwrap();
let storage: Box<dyn Storage> = Box::new(LocalStorage::new_from_path(tmp_file.path()));
storage.store("Java is awesome".as_bytes()).await.unwrap();

tmp_file.as_file().sync_all().unwrap();

// corrupt the file
tmp_file.seek(SeekFrom::Start(0)).unwrap();
tmp_file.write_all("Raft".as_bytes()).unwrap();

tmp_file.as_file().sync_all().unwrap();

let result = storage.turned_malicious().await;
assert!(result.is_err());
}

#[tokio::test]
async fn test_turned_malicious_happy_case() {
let tmp_file = NamedTempFile::new().unwrap();
let storage: Box<dyn Storage> = Box::new(LocalStorage::new_from_path(tmp_file.path()));
storage.store("Java is awesome".as_bytes()).await.unwrap();

tmp_file.as_file().sync_all().unwrap();

let result = storage.turned_malicious().await;
assert!(result.is_ok());
}
}

0 comments on commit 8530aa7

Please sign in to comment.