From 3b8b260b6505bbb2a1a7aac1740eca45d28192c5 Mon Sep 17 00:00:00 2001
From: Paul Campbell <pcampbell@kemitix.net>
Date: Sun, 1 Dec 2024 20:58:05 +0000
Subject: [PATCH] feat(net): add user_agent helper to MockNet

---
 src/net/system.rs | 12 ++++++++++++
 tests/net.rs      | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/src/net/system.rs b/src/net/system.rs
index fa56fdf..8e12ed4 100644
--- a/src/net/system.rs
+++ b/src/net/system.rs
@@ -434,6 +434,12 @@ impl<'net> ReqBuilder<'net> {
         self
     }
 
+    /// Sets the User-Agent header to be used by this client.
+    #[must_use]
+    pub fn user_agent(self, user_agent: impl Into<String>) -> Self {
+        self.header(http::header::USER_AGENT.to_string(), user_agent)
+    }
+
     /// Add query parameter
     #[must_use]
     pub fn query(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
@@ -756,6 +762,12 @@ impl<'net> WhenRequest<'net, WhenBuildRequest> {
         self
     }
 
+    /// Specifies user agent the mock will match against.
+    #[must_use]
+    pub fn user_agent(self, agent: impl Into<String>) -> Self {
+        self.header(http::header::USER_AGENT.to_string(), agent)
+    }
+
     /// Specifies the body that the mock will match against.
     ///
     /// Any request that does not have this body will not match the mock.
diff --git a/tests/net.rs b/tests/net.rs
index d98215a..a674f2e 100644
--- a/tests/net.rs
+++ b/tests/net.rs
@@ -531,3 +531,40 @@ async fn test_get_with_duplicate_query_keys() {
         "key:value-3,value-4"
     );
 }
+
+#[tokio::test]
+async fn test_get_with_user_agent() {
+    //given
+    let mock_net = kxio::net::mock();
+    let url = "https://www.example.com/path";
+
+    mock_net
+        .on()
+        .get(url)
+        .user_agent("007")
+        .respond(StatusCode::OK)
+        .mock()
+        .expect("mock");
+    mock_net
+        .on()
+        .get(url)
+        .user_agent("orange")
+        .respond(StatusCode::FORBIDDEN)
+        .mock()
+        .expect("mock");
+    let net = Net::from(mock_net);
+
+    //when
+    let invalid = net.get(url).user_agent("orange").send().await;
+    let valid = net
+        .get(url)
+        .user_agent("007")
+        .send()
+        .await
+        .expect("valid response");
+
+    //then
+    let_assert!(Err(Error::ResponseError { response }) = invalid);
+    assert_eq!(response.status(), StatusCode::FORBIDDEN);
+    assert_eq!(valid.status(), StatusCode::OK);
+}