From 5df74903228a7abb305eba4743b9015b9ba49c63 Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil.hafeez@gmail.com>
Date: Thu, 15 May 2025 23:02:51 -0700
Subject: [PATCH] add tests for parse response

---
 .../brightstaff/src/router/router_model_v1.rs | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/crates/brightstaff/src/router/router_model_v1.rs b/crates/brightstaff/src/router/router_model_v1.rs
index fdff34c2..7bd1a987 100644
--- a/crates/brightstaff/src/router/router_model_v1.rs
+++ b/crates/brightstaff/src/router/router_model_v1.rs
@@ -198,3 +198,46 @@ user: "seattle"
         assert_eq!(expected_prompt, prompt);
     }
 }
+
+#[test]
+fn test_parse_response() {
+    let router = RouterModelV1::new(
+        "route1: description1\nroute2: description2".to_string(),
+        "test-model".to_string(),
+    );
+
+    // Case 1: Valid JSON with non-empty route
+    let input = r#"{"route": "route1"}"#;
+    let result = router.parse_response(input).unwrap();
+    assert_eq!(result, Some("route1".to_string()));
+
+    // Case 2: Valid JSON with empty route
+    let input = r#"{"route": ""}"#;
+    let result = router.parse_response(input).unwrap();
+    assert_eq!(result, None);
+
+    // Case 3: Valid JSON with null route
+    let input = r#"{"route": null}"#;
+    let result = router.parse_response(input).unwrap();
+    assert_eq!(result, None);
+
+    // Case 4: JSON missing route field
+    let input = r#"{}"#;
+    let result = router.parse_response(input).unwrap();
+    assert_eq!(result, None);
+
+    // Case 5: Malformed JSON
+    let input = r#"{"route": "route1""#; // missing closing }
+    let result = router.parse_response(input);
+    assert!(result.is_err());
+
+    // Case 6: Single quotes and \n in JSON
+    let input = "{'route': 'route2'}\\n";
+    let result = router.parse_response(input).unwrap();
+    assert_eq!(result, Some("route2".to_string()));
+
+    // Case 7: Code block marker
+    let input = "```json\n{\"route\": \"route1\"}\n```";
+    let result = router.parse_response(input).unwrap();
+    assert_eq!(result, Some("route1".to_string()));
+}