chore: more tests, never enough

dreadnode · Dec 10, 2024 · 15a99b3 · 15a99b3
1 parent f468f72
commit 15a99b3
Showing 1 changed file with 43 additions and 0 deletions.
diff --git a/src/agent/generator/history.rs b/src/agent/generator/history.rs
@@ -315,4 +315,47 @@ mod tests {
         let history = ChatHistory::create(conv, ConversationWindow::LastN(4));
         assert_eq!(history.history, expected);
     }
+
+    #[test]
+    fn test_last_n_strategy_with_small_conv() {
+        let conv = vec![
+            Message::Agent("test1".to_string(), None),
+            Message::Feedback("feedback1".to_string(), None),
+        ];
+
+        let expected = vec![
+            Message::Agent("test1".to_string(), None),
+            Message::Feedback("feedback1".to_string(), None),
+        ];
+
+        let history = ChatHistory::create(conv, ConversationWindow::LastN(10));
+        assert_eq!(history.history, expected);
+    }
+
+    #[test]
+    fn test_last_n_strategy_with_just_enough() {
+        let conv = vec![
+            Message::Agent("test1".to_string(), None),
+            Message::Feedback("feedback1".to_string(), None),
+            Message::Agent("test2".to_string(), None),
+            Message::Feedback("feedback2".to_string(), None),
+            Message::Agent("test3".to_string(), None),
+            Message::Feedback("feedback3".to_string(), None),
+            Message::Agent("test4".to_string(), None),
+            Message::Feedback("feedback4".to_string(), None),
+        ];
+
+        let expected = vec![
+            Message::Feedback("feedback1".to_string(), None),
+            Message::Agent("test2".to_string(), None),
+            Message::Feedback("feedback2".to_string(), None),
+            Message::Agent("test3".to_string(), None),
+            Message::Feedback("feedback3".to_string(), None),
+            Message::Agent("test4".to_string(), None),
+            Message::Feedback("feedback4".to_string(), None),
+        ];
+
+        let history = ChatHistory::create(conv, ConversationWindow::LastN(7));
+        assert_eq!(history.history, expected);
+    }
 }