From ffe6193a6ecb63b9de4e9bdbb4d2e7397ecae5c6 Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil@katanemo.com>
Date: Tue, 22 Oct 2024 12:42:44 -0700
Subject: [PATCH] pretty_eq

---
 crates/Cargo.lock                          | 1 +
 crates/prompt_gateway/Cargo.toml           | 1 +
 crates/prompt_gateway/src/hallucination.rs | 3 +++
 3 files changed, 5 insertions(+)

diff --git a/crates/Cargo.lock b/crates/Cargo.lock
index 4a76be14..a9c06a49 100644
--- a/crates/Cargo.lock
+++ b/crates/Cargo.lock
@@ -1120,6 +1120,7 @@ dependencies = [
  "http",
  "log",
  "md5",
+ "pretty_assertions",
  "proxy-wasm",
  "proxy-wasm-test-framework",
  "rand",
diff --git a/crates/prompt_gateway/Cargo.toml b/crates/prompt_gateway/Cargo.toml
index 29d385b7..e8a166f8 100644
--- a/crates/prompt_gateway/Cargo.toml
+++ b/crates/prompt_gateway/Cargo.toml
@@ -26,3 +26,4 @@ sha2 = "0.10.8"
 [dev-dependencies]
 proxy-wasm-test-framework = { git = "https://github.com/katanemo/test-framework.git", branch = "new" }
 serial_test = "3.1.1"
+pretty_assertions = "1.4.1"
diff --git a/crates/prompt_gateway/src/hallucination.rs b/crates/prompt_gateway/src/hallucination.rs
index a17c6f60..62b119ac 100644
--- a/crates/prompt_gateway/src/hallucination.rs
+++ b/crates/prompt_gateway/src/hallucination.rs
@@ -40,6 +40,7 @@ pub fn extract_messages_for_hallucination(messages: &Vec<Message>) -> Vec<String
 
 #[cfg(test)]
 mod test {
+    use pretty_assertions::assert_eq;
     use common::common_types::open_ai::Message;
 
     use super::extract_messages_for_hallucination;
@@ -157,5 +158,7 @@ mod test {
         let messages_for_halluncination = extract_messages_for_hallucination(&messages);
         println!("{:?}", messages_for_halluncination);
         assert_eq!(messages_for_halluncination.len(), 3);
+        assert_eq!(["tell me about the weather", "Seattle", "7 days"], messages_for_halluncination.as_slice());
     }
+
 }