Fixed some tests

3 files changed, 77 insertions, 4 deletions

diff --git a/crates/utils/tcp_connection/tcp_connection_test/src/lib.rs b/crates/utils/tcp_connection/tcp_connection_test/src/lib.rs
index 2774458..f0eb66e 100644
--- a/crates/utils/tcp_connection/tcp_connection_test/src/lib.rs
+++ b/crates/utils/tcp_connection/tcp_connection_test/src/lib.rs
@@ -6,3 +6,6 @@ pub mod test_connection;
 
 #[cfg(test)]
 pub mod test_challenge;
+
+#[cfg(test)]
+pub mod test_file_transfer;
diff --git a/crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs b/crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs
index 723bf52..05e47ad 100644
--- a/crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs
+++ b/crates/utils/tcp_connection/tcp_connection_test/src/test_challenge.rs
@@ -18,9 +18,50 @@ impl ClientHandle<ExampleChallengeServerHandle> for ExampleChallengeClientHandle
         mut instance: ConnectionInstance,
     ) -> impl std::future::Future<Output = ()> + Send + Sync {
         async move {
-            let key = current_dir().unwrap().join("res").join("test_key");
+            // Accept challenge with correct key
+            let key = current_dir()
+                .unwrap()
+                .join("res")
+                .join("key")
+                .join("test_key_private.pem");
             let result = instance.accept_challenge(key, "test_key").await.unwrap();
+
+            // Sent success
+            assert_eq!(true, result);
+            let response = instance.read_text().await.unwrap();
+
+            // Verify success
+            assert_eq!("OK", response);
+
+            // Accept challenge with wrong key
+            let key = current_dir()
+                .unwrap()
+                .join("res")
+                .join("key")
+                .join("wrong_key_private.pem");
+            let result = instance.accept_challenge(key, "test_key").await.unwrap();
+
+            // Sent success
+            assert_eq!(true, result);
+            let response = instance.read_text().await.unwrap();
+
+            // Verify fail
+            assert_eq!("ERROR", response);
+
+            // Accept challenge with wrong name
+            let key = current_dir()
+                .unwrap()
+                .join("res")
+                .join("key")
+                .join("test_key_private.pem");
+            let result = instance.accept_challenge(key, "test_key__").await.unwrap();
+
+            // Sent success
             assert_eq!(true, result);
+            let response = instance.read_text().await.unwrap();
+
+            // Verify fail
+            assert_eq!("ERROR", response);
         }
     }
 }
@@ -32,9 +73,38 @@ impl ServerHandle<ExampleChallengeClientHandle> for ExampleChallengeServerHandle
         mut instance: ConnectionInstance,
     ) -> impl std::future::Future<Output = ()> + Send + Sync {
         async move {
-            let key_dir = current_dir().unwrap().join("res");
+            // Challenge with correct key
+            let key_dir = current_dir().unwrap().join("res").join("key");
             let result = instance.challenge(key_dir).await.unwrap();
             assert_eq!(true, result);
+
+            // Send response
+            instance
+                .write_text(if result { "OK" } else { "ERROR" })
+                .await
+                .unwrap();
+
+            // Challenge again
+            let key_dir = current_dir().unwrap().join("res").join("key");
+            let result = instance.challenge(key_dir).await.unwrap();
+            assert_eq!(false, result);
+
+            // Send response
+            instance
+                .write_text(if result { "OK" } else { "ERROR" })
+                .await
+                .unwrap();
+
+            // Challenge again
+            let key_dir = current_dir().unwrap().join("res").join("key");
+            let result = instance.challenge(key_dir).await.unwrap();
+            assert_eq!(false, result);
+
+            // Send response
+            instance
+                .write_text(if result { "OK" } else { "ERROR" })
+                .await
+                .unwrap();
         }
     }
 }
diff --git a/crates/utils/tcp_connection/tcp_connection_test/src/test_connection.rs b/crates/utils/tcp_connection/tcp_connection_test/src/test_connection.rs
index f35fd80..bbea92f 100644
--- a/crates/utils/tcp_connection/tcp_connection_test/src/test_connection.rs
+++ b/crates/utils/tcp_connection/tcp_connection_test/src/test_connection.rs
@@ -20,7 +20,7 @@ impl ClientHandle<ExampleServerHandle> for ExampleClientHandle {
                 panic!("Write text failed!");
             };
             // Read msg
-            let Ok(result) = instance.read_text(512 as u32).await else {
+            let Ok(result) = instance.read_text().await else {
                 return;
             };
             assert_eq!("Hello Peter!", result);
@@ -36,7 +36,7 @@ impl ServerHandle<ExampleClientHandle> for ExampleServerHandle {
     ) -> impl std::future::Future<Output = ()> + Send + Sync {
         async move {
             // Read name
-            let Ok(name) = instance.read_text(512 as u32).await else {
+            let Ok(name) = instance.read_text().await else {
                 return;
             };
             // Write msg