refactor: Clean up fig_test_utils (#440)

Author: brandonskiser

crates/fig_install/src/linux.rs

@@ -319,6 +319,7 @@ mod tests {
     use fig_test_utils::http::Method;
     use fig_util::CLI_BINARY_NAME;
     use fig_util::directories::home_dir;
+    use hex::ToHex;
 
     use super::*;
 
@@ -410,7 +411,7 @@ mod tests {
             UpdatePackage {
                 version: semver::Version::from_str(test_version).unwrap(),
                 download_url: Url::from_str(&format!("http://{}{}", test_server_addr, test_download_path)).unwrap(),
-                sha256: fig_test_utils::sha256(test_file.as_bytes()),
+                sha256: ring::digest::digest(&ring::digest::SHA256, test_file.as_bytes()).encode_hex(),
                 size: 0, // size not checked
                 cli_path: None,
             },

crates/fig_test_utils/src/lib.rs

@@ -1,10 +1,4 @@
 pub mod server;
 
-use hex::ToHex;
 pub use http;
 pub use server::*;
-
-/// Encodes the byte slice as a lower-case sha256 hex string.
-pub fn sha256(data: &[u8]) -> String {
-    ring::digest::digest(&ring::digest::SHA256, data).encode_hex()
-}

crates/fig_test_utils/src/server.rs

@@ -4,6 +4,7 @@ use std::net::SocketAddr;
 use std::pin::Pin;
 
 use bytes::Bytes;
+pub use http::Method;
 use http_body_util::Full;
 use hyper::body::Incoming;
 use hyper::service::Service;
@@ -19,10 +20,34 @@ use tokio::net::{
 use tokio::select;
 use tokio_util::sync::CancellationToken;
 
+/// A local HTTP server that can be used for returning mock responses to HTTP requests.
+///
+/// # Examples
+///
+/// ```rust
+/// # use fig_test_utils::server::*;
+/// # async fn run() -> Result<(), reqwest::Error> {
+/// // Hosting a local server that responds to GET requests for "/my-file" with the
+/// // body "some text".
+/// let test_path = String::from("/my-file");
+/// let mock_body = String::from("some text");
+/// let test_server_addr = TestServer::new()
+///     .await
+///     .with_mock_response(Method::GET, test_path.clone(), mock_body.clone())
+///     .spawn_listener();
+///
+/// let body = reqwest::get(format!("http://{}{}", &test_server_addr, &test_path))
+///     .await?
+///     .text()
+///     .await?;
+///
+/// assert_eq!(body, mock_body);
+/// # Ok(())
+/// # }
+/// ```
 #[derive(Debug)]
 pub struct TestServer {
     listener: TcpListener,
-    cancellation_token: Option<CancellationToken>,
     mock_responses: HashMap<(http::Method, String), String>,
 }
 
@@ -33,7 +58,6 @@ impl TestServer {
             .expect("failed to bind socket to local host");
         Self {
             listener,
-            cancellation_token: None,
             mock_responses: HashMap::default(),
         }
     }
@@ -43,14 +67,14 @@ impl TestServer {
         self
     }
 
-    pub fn spawn_listener(mut self) -> SocketAddr {
-        let addr = self
+    /// Spawns a new task for accepting requests, returning the address of the listening socket.
+    pub fn spawn_listener(self) -> TestAddress {
+        let address = self
             .listener
             .local_addr()
             .expect("listener should be bound to an address");
-        let token = CancellationToken::new();
-        let token_clone = token.clone();
-        self.cancellation_token = Some(token);
+        let cancellation_token = CancellationToken::new();
+        let token_clone = cancellation_token.clone();
         tokio::task::spawn(async move {
             loop {
                 select! {
@@ -64,7 +88,10 @@ impl TestServer {
                 }
             }
         });
-        addr
+        TestAddress {
+            address,
+            cancellation_token,
+        }
     }
 
     async fn handle_request(&self, stream: TcpStream) {
@@ -76,11 +103,27 @@ impl TestServer {
     }
 }
 
-impl Drop for TestServer {
+#[derive(Debug)]
+pub struct TestAddress {
+    address: SocketAddr,
+    cancellation_token: CancellationToken,
+}
+
+impl TestAddress {
+    pub fn address(&self) -> SocketAddr {
+        self.address
+    }
+}
+
+impl std::fmt::Display for TestAddress {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.address)
+    }
+}
+
+impl Drop for TestAddress {
     fn drop(&mut self) {
-        if let Some(token) = &self.cancellation_token {
-            token.cancel();
-        }
+        self.cancellation_token.cancel();
     }
 }
 
@@ -100,3 +143,34 @@ impl Service<Request<Incoming>> for TestServer {
         Box::pin(async move { Ok(Response::builder().status(200).body(result.into()).unwrap()) })
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_server_mock_and_drop() {
+        let test_path = String::from("/test-path");
+        let test_response = String::from("test body");
+        let test_server_addr = TestServer::new()
+            .await
+            .with_mock_response(Method::GET, test_path.clone(), test_response.clone())
+            .spawn_listener();
+
+        let response = reqwest::get(format!("http://{}{}", &test_server_addr, &test_path))
+            .await
+            .unwrap()
+            .text()
+            .await
+            .unwrap();
+        assert_eq!(response, test_response);
+
+        // Test that dropping TestAddress stops the server
+        let addr = test_server_addr.address();
+        std::mem::drop(test_server_addr);
+        // wait for the task to complete
+        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+        let response = reqwest::get(format!("http://{}{}", &addr, &test_path)).await;
+        assert!(response.is_err());
+    }
+}