test(tedge-p11-server): simplify unit tests

Signed-off-by: Marcel Guzik <marcel.guzik@cumulocity.com>

Author: Bravo555

crates/extensions/tedge-p11-server/Cargo.toml

@@ -29,6 +29,7 @@ tokio = { workspace = true, features = [
     "rt-multi-thread",
     "signal",
     "net",
+    "io-util",
 ] }
 toml.workspace = true
 tracing.workspace = true

crates/extensions/tedge-p11-server/src/proxy/server.rs

@@ -139,9 +139,10 @@ mod tests {
     use crate::pkcs11;
     use crate::proxy::frame::Frame;
     use crate::service::*;
-    use std::io::Read;
     use std::os::unix::net::UnixStream;
     use std::time::Duration;
+    use tokio::io::AsyncReadExt as _;
+    use tokio::io::AsyncWriteExt as _;
 
     const SCHEME: pkcs11::SigScheme = pkcs11::SigScheme::EcdsaNistp256Sha256;
     const SIGNATURE: [u8; 2] = [0x21, 0x37];
@@ -180,15 +181,7 @@ mod tests {
     /// connection, framing, serialization, but not PKCS#11 layer itself.
     #[tokio::test]
     async fn server_works_with_client() {
-        let service = TestSigningService;
-        let server = TedgeP11Server::new(service).unwrap();
-        let tmpdir = tempfile::tempdir().unwrap();
-        let socket_path = tmpdir.path().join("test_socket.sock");
-        let listener = UnixListener::bind(&socket_path).unwrap();
-
-        tokio::spawn(async move { server.serve(listener).await });
-        // wait until the server calls accept()
-        tokio::time::sleep(Duration::from_millis(2)).await;
+        let (socket_path, _s) = setup_server().await;
 
         tokio::task::spawn_blocking(move || {
             let client = TedgeP11Client::with_ready_check(socket_path.into());
@@ -204,53 +197,29 @@ mod tests {
 
     #[tokio::test]
     async fn server_responds_with_error_to_invalid_request() {
-        let service = TestSigningService;
-        let server = TedgeP11Server::new(service).unwrap();
-        let tmpdir = tempfile::tempdir().unwrap();
-        let socket_path = tmpdir.path().join("test_socket.sock");
-        let listener = UnixListener::bind(&socket_path).unwrap();
-
-        tokio::spawn(async move { server.serve(listener).await });
-        // wait until the server calls accept()
-        tokio::time::sleep(Duration::from_millis(2)).await;
+        let (socket_path, _s) = setup_server().await;
 
         let response = tokio::task::spawn_blocking(move || {
-            let mut client_connection = Connection::new(UnixStream::connect(socket_path).unwrap());
-            client_connection
+            let mut client = Connection::new(UnixStream::connect(socket_path).unwrap());
+            client
                 .write_frame(&Frame1::SignResponse(SignResponse(vec![])))
                 .unwrap();
-            client_connection.read_frame().unwrap()
+            client.read_frame().unwrap()
         })
         .await
         .unwrap();
-        assert!(matches!(response, Frame1::Error(_)));
+        let Frame1::Error(ProtocolError(err_msg)) = response else {
+            panic!("should be error");
+        };
+        assert_eq!(err_msg.as_str(), "invalid request");
     }
 
     #[tokio::test]
     async fn server_responds_with_error_to_garbage() {
-        use std::io::Write as _;
-
-        let service = TestSigningService;
-        let server = TedgeP11Server::new(service).unwrap();
-        let tmpdir = tempfile::tempdir().unwrap();
-        let socket_path = tmpdir.path().join("test_socket.sock");
-        let listener = UnixListener::bind(&socket_path).unwrap();
-
-        tokio::spawn(async move { server.serve(listener).await });
-        // wait until the server calls accept()
-        tokio::time::sleep(Duration::from_millis(2)).await;
+        let (mut client, _s) = setup_test().await;
 
-        // the reader should exit
-        let response = tokio::task::spawn_blocking(move || {
-            let mut stream = UnixStream::connect(socket_path).unwrap();
-            write!(stream, "garbage").unwrap();
-            stream.shutdown(std::net::Shutdown::Write).unwrap();
-            let mut response = Vec::new();
-            stream.read_to_end(&mut response).unwrap();
-            response
-        })
-        .await
-        .unwrap();
+        client.write_and_close("garbage".as_bytes()).await;
+        let response = client.read().await;
 
         let response: Frame = postcard::from_bytes(&response).unwrap();
         let Frame::Version1(Frame1::Error(ProtocolError(err_msg))) = response else {
@@ -262,4 +231,48 @@ mod tests {
             "read: Failed to parse the received frame: Serde Deserialization Error"
         );
     }
+
+    async fn setup_test() -> (TestClient, tokio::task::JoinHandle<anyhow::Result<()>>) {
+        let (socket_path, server) = setup_server().await;
+
+        let client_socket = TestClient(tokio::net::UnixStream::connect(socket_path).await.unwrap());
+        (client_socket, server)
+    }
+
+    async fn setup_server() -> (
+        std::path::PathBuf,
+        tokio::task::JoinHandle<Result<(), anyhow::Error>>,
+    ) {
+        let tmpdir = tempfile::tempdir().unwrap();
+        let server = TedgeP11Server::new(TestSigningService).unwrap();
+        let socket_path = tmpdir.path().join("test_socket.sock");
+        let listener = UnixListener::bind(&socket_path).unwrap();
+
+        let server = tokio::spawn(async move {
+            let _tmpdir = tmpdir; // keep tmpdir alive until the server is done
+            server.serve(listener).await.unwrap();
+            Ok(())
+        });
+
+        // wait until the server calls accept()
+        tokio::time::sleep(Duration::from_millis(100)).await;
+
+        (socket_path, server)
+    }
+
+    struct TestClient(tokio::net::UnixStream);
+
+    impl TestClient {
+        async fn write_and_close(&mut self, bytes: &[u8]) {
+            self.0.write_all(bytes).await.unwrap();
+            self.0.flush().await.unwrap();
+            self.0.shutdown().await.unwrap();
+        }
+
+        async fn read(&mut self) -> Vec<u8> {
+            let mut bytes = Vec::new();
+            self.0.read_to_end(&mut bytes).await.unwrap();
+            bytes
+        }
+    }
 }