From ffc74ab694e643ed594e73efac73962b69ee61e6 Mon Sep 17 00:00:00 2001
From: Jian Cao <jiancao0421@meta.com>
Date: Mon, 24 Oct 2022 18:02:08 -0700
Subject: [PATCH] Add unit test for tls and streaming

Summary:
# What
* Add unit test for tls. Use rcgen to create the cert and pem and run from_dir()
* Add unit test for send_data. Just send the payload and see whether it can finish or not

# Why
* need to improve code coverage

Differential Revision: D40660872

fbshipit-source-id: 10ff8f269689fbf692f3bb2a1cf7351190b729f2
---
 protocol-rpc/Cargo.toml             |  2 +-
 protocol-rpc/src/connect/tls.rs     | 44 +++++++++++++++++++++++++++++
 protocol-rpc/src/proto/streaming.rs | 39 +++++++++++++++++++++++++
 3 files changed, 84 insertions(+), 1 deletion(-)

diff --git a/protocol-rpc/Cargo.toml b/protocol-rpc/Cargo.toml
index cba4030..489623e 100644
--- a/protocol-rpc/Cargo.toml
+++ b/protocol-rpc/Cargo.toml
@@ -92,7 +92,7 @@ itertools = "0.9.0"
 tempfile = "3.2.0"
 num-bigint = { version = "0.4", features = ["rand"] }
 num-traits = "0.2"
-
+rcgen = "0.10.0"
 
 [build-dependencies]
 tonic-build = { version = "0.7.2" }
diff --git a/protocol-rpc/src/connect/tls.rs b/protocol-rpc/src/connect/tls.rs
index cd31f14..313cee9 100644
--- a/protocol-rpc/src/connect/tls.rs
+++ b/protocol-rpc/src/connect/tls.rs
@@ -147,4 +147,48 @@ mod tests {
             "https"
         );
     }
+
+    #[test]
+    fn test_host_into_uri() {
+        let _ = host_into_uri("foo.bar:10009/path?1234", false);
+    }
+
+    #[tokio::test]
+    async fn test_from_dir() {
+        use std::fs::File;
+        use std::io::Write;
+
+        use tempfile::tempdir;
+
+        // Create a directory inside of `std::env::temp_dir()`.
+        let dir = tempdir().unwrap();
+        use rcgen::*;
+        let subject_alt_names: &[_] = &["hello.world.example".to_string(), "localhost".to_string()];
+        let ca_subject_alt_names: &[_] = &["ca.world.example".to_string(), "localhost".to_string()];
+
+        let client_cert = generate_simple_self_signed(subject_alt_names).unwrap();
+        let ca_cert = generate_simple_self_signed(ca_subject_alt_names).unwrap();
+        let client_pem = client_cert.serialize_pem().unwrap();
+        let client_key = client_cert.serialize_private_key_pem();
+        let ca_pem = ca_cert.serialize_pem().unwrap();
+
+        let file_path_ca_pem = dir.path().join("ca.pem");
+        let mut file_ca_pem = File::create(file_path_ca_pem).unwrap();
+        file_ca_pem.write_all(ca_pem.as_bytes()).unwrap();
+
+        let file_path_client_pem = dir.path().join("client.pem");
+        let mut file_client_pem = File::create(file_path_client_pem).unwrap();
+        file_client_pem.write_all(client_pem.as_bytes()).unwrap();
+
+        let file_path_client_key = dir.path().join("client.key");
+        let mut file_client_key = File::create(file_path_client_key).unwrap();
+        file_client_key.write_all(client_key.as_bytes()).unwrap();
+
+        let _ = TlsContext::from_dir(dir.path(), false);
+
+        drop(file_ca_pem);
+        drop(file_client_pem);
+        drop(file_client_key);
+        dir.close().unwrap();
+    }
 }
diff --git a/protocol-rpc/src/proto/streaming.rs b/protocol-rpc/src/proto/streaming.rs
index ff23694..a756027 100644
--- a/protocol-rpc/src/proto/streaming.rs
+++ b/protocol-rpc/src/proto/streaming.rs
@@ -54,3 +54,42 @@ pub fn write_to_stream(payload: TPayload) -> Response<TPayloadStream> {
     });
     Response::new(ReceiverStream::new(rx))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_send_data() {
+        let data = vec![
+            ByteBuffer {
+                buffer: vec![
+                    200, 135, 56, 19, 5, 207, 16, 147, 198, 229, 224, 111, 97, 119, 247, 238, 48,
+                    209, 55, 188, 30, 178, 53, 4, 110, 27, 182, 220, 156, 57, 53, 63,
+                ],
+            },
+            ByteBuffer {
+                buffer: vec![
+                    102, 237, 233, 208, 207, 235, 165, 5, 177, 27, 168, 233, 239, 69, 163, 80, 155,
+                    2, 85, 192, 182, 25, 20, 189, 118, 5, 225, 153, 13, 254, 201, 40,
+                ],
+            },
+            ByteBuffer {
+                buffer: vec![
+                    48, 54, 39, 197, 69, 34, 214, 167, 225, 117, 64, 223, 51, 164, 33, 208, 18,
+                    108, 38, 248, 215, 189, 94, 180, 82, 105, 196, 43, 189, 2, 220, 6,
+                ],
+            },
+            ByteBuffer {
+                buffer: vec![
+                    228, 188, 46, 30, 21, 100, 156, 96, 162, 185, 103, 149, 89, 159, 81, 67, 119,
+                    112, 0, 174, 99, 188, 74, 7, 13, 236, 98, 48, 50, 145, 156, 50,
+                ],
+            },
+        ];
+
+        let p = Payload::from(&data);
+        let tp = TPayload::from(&p);
+        send_data(tp);
+    }
+}