Mock DDG in unit tests (#1076)

griptape-ai · Aug 19, 2024 · 717753c · 717753c
1 parent 3aaeb6e
commit 717753c
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/tests/unit/drivers/web_search/test_duck_duck_go_web_search_driver.py b/tests/unit/drivers/web_search/test_duck_duck_go_web_search_driver.py
@@ -13,14 +13,21 @@ def driver(self, mocker):
             {"title": "foo", "href": "bar", "body": "baz"},
             {"title": "foo2", "href": "bar2", "body": "baz2"},
         ]
-
-        mocker.patch("duckduckgo_search.DDGS.text", return_value=mock_response)
-
+        mock_ddg = mocker.Mock(
+            text=lambda *args, **kwargs: mock_response,
+        )
+        mocker.patch("duckduckgo_search.DDGS", return_value=mock_ddg)
         return DuckDuckGoWebSearchDriver()
 
     @pytest.fixture()
     def driver_with_error(self, mocker):
-        mocker.patch("duckduckgo_search.DDGS.text", side_effect=Exception("test_error"))
+        def error(*args, **kwargs):
+            raise Exception("test_error")
+
+        mock_ddg = mocker.Mock(
+            text=error,
+        )
+        mocker.patch("duckduckgo_search.DDGS", return_value=mock_ddg)
 
         return DuckDuckGoWebSearchDriver()