From 92b19ed1fb5cfd20c84c00d8bbb3c905f6e96bdb Mon Sep 17 00:00:00 2001
From: Graham Neubig <neubig@gmail.com>
Date: Wed, 14 Aug 2024 04:20:30 -0400
Subject: [PATCH] Add unit tests for MemoryCondenser in test_condenser.py
 (#3379)

* Add unit tests for MemoryCondenser in test_condenser.py

* Formatting

* Fix formatting etc

---------

Co-authored-by: opendevin <opendevin@all-hands.dev>
---
 tests/unit/test_condenser.py | 44 ++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)
 create mode 100644 tests/unit/test_condenser.py

diff --git a/tests/unit/test_condenser.py b/tests/unit/test_condenser.py
new file mode 100644
index 0000000000..39a882ce4e
--- /dev/null
+++ b/tests/unit/test_condenser.py
@@ -0,0 +1,44 @@
+from unittest.mock import Mock, patch
+
+import pytest
+
+from opendevin.core.exceptions import LLMResponseError
+from opendevin.llm.llm import LLM
+from opendevin.memory.condenser import MemoryCondenser
+
+
+@pytest.fixture
+def memory_condenser():
+    return MemoryCondenser()
+
+
+@pytest.fixture
+def mock_llm():
+    return Mock(spec=LLM)
+
+
+def test_condense_success(memory_condenser, mock_llm):
+    mock_llm.completion.return_value = {
+        'choices': [{'message': {'content': 'Condensed memory'}}]
+    }
+    result = memory_condenser.condense('Summarize this', mock_llm)
+    assert result == 'Condensed memory'
+    mock_llm.completion.assert_called_once_with(
+        messages=[{'content': 'Summarize this', 'role': 'user'}]
+    )
+
+
+def test_condense_exception(memory_condenser, mock_llm):
+    mock_llm.completion.side_effect = LLMResponseError('LLM error')
+    with pytest.raises(LLMResponseError, match='LLM error'):
+        memory_condenser.condense('Summarize this', mock_llm)
+
+
+@patch('opendevin.memory.condenser.logger')
+def test_condense_logs_error(mock_logger, memory_condenser, mock_llm):
+    mock_llm.completion.side_effect = LLMResponseError('LLM error')
+    with pytest.raises(LLMResponseError):
+        memory_condenser.condense('Summarize this', mock_llm)
+    mock_logger.error.assert_called_once_with(
+        'Error condensing thoughts: %s', 'LLM error', exc_info=False
+    )