From 83d430a310b04bc0a2ea49612fce0de70558d63a Mon Sep 17 00:00:00 2001
From: Kurt Heiden <kurt.w.heiden@gmail.com>
Date: Mon, 28 Oct 2024 21:59:36 -0600
Subject: [PATCH] Update test_ollama.py

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
---
 test_ollama.py | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/test_ollama.py b/test_ollama.py
index 76a7cc2..f4b721c 100644
--- a/test_ollama.py
+++ b/test_ollama.py
@@ -48,17 +48,21 @@ class TestOllama(unittest.TestCase):
     self.assertIsInstance(result, sm.models.Message)
 
   def test_structure_response(self):
-    class Poem(BaseModel):
-        title: str
-        content: str
-    with self.assertRaises(NotImplementedError):
-      data_obj = sm.generate_data(
-        prompt="Write a poem about love",
-        llm_provider="ollama",
-        llm_model="llama3.2",
-        response_model=Poem)
-      self.assertIsNotNone(data_obj)
-      self.assertIsInstance(data_obj, Poem)
+      class Poem(BaseModel):
+          title: str
+          content: str
+      # Test for NotImplementedError
+      with self.assertRaises(NotImplementedError):
+          sm.generate_data(
+              prompt="Write a poem about love",
+              llm_provider="ollama",
+              llm_model="llama3.2",
+              response_model=Poem)
+      
+      # If implementation is added later, uncomment and modify:
+      # data_obj = sm.generate_data(...)
+      # self.assertIsNotNone(data_obj)
+      # self.assertIsInstance(data_obj, Poem)
 
 
 if __name__ == '__main__':