Fix TypeError: 'Mock' object is not iterable

Author: yvan-sraka

tests/test_outline.py

@@ -1,11 +1,18 @@
-from unittest.mock import Mock, patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 from pydantic import BaseModel
 
 from outlines import Outline
 
 
+class IterableMock(MagicMock):
+    def __getattr__(self, name):
+        result = MagicMock()
+        result.__iter__.return_value = iter([])
+        return result
+
+
 class OutputModel(BaseModel):
     result: int
 
@@ -15,8 +22,8 @@ def template(a: int) -> str:
 
 
 def test_outline():
-    mock_model = Mock()
-    mock_generator = Mock()
+    mock_model = IterableMock()
+    mock_generator = MagicMock()
     mock_generator.return_value = '{"result": 6}'
     with patch("outlines.generate.json", return_value=mock_generator):
         outline_instance = Outline(mock_model, template, OutputModel)
@@ -26,8 +33,8 @@ def test_outline():
 
 
 def test_outline_with_json_schema():
-    mock_model = Mock()
-    mock_generator = Mock()
+    mock_model = IterableMock()
+    mock_generator = MagicMock()
     mock_generator.return_value = '{"result": 6}'
     with patch("outlines.generate.json", return_value=mock_generator):
         outline_instance = Outline(
@@ -40,14 +47,14 @@ def test_outline_with_json_schema():
 
 
 def test_invalid_output_type():
-    mock_model = Mock()
+    mock_model = IterableMock()
     with pytest.raises(TypeError):
         Outline(mock_model, template, int)
 
 
 def test_invalid_json_response():
-    mock_model = Mock()
-    mock_generator = Mock()
+    mock_model = IterableMock()
+    mock_generator = MagicMock()
     mock_generator.return_value = "invalid json"
     with patch("outlines.generate.json", return_value=mock_generator):
         outline_instance = Outline(mock_model, template, OutputModel)
@@ -56,7 +63,7 @@ def test_invalid_json_response():
 
 
 def test_invalid_json_schema():
-    mock_model = Mock()
+    mock_model = IterableMock()
     invalid_json_schema = (
         '{"type": "object", "properties": {"result": {"type": "invalid_type"}}}'
     )