diff --git a/test_ollama.py b/test_ollama.py
index f4b721c..0a5b560 100644
--- a/test_ollama.py
+++ b/test_ollama.py
@@ -1,4 +1,6 @@
+import os
 import unittest
+from unittest import mock
 import simplemind as sm
 from pydantic import BaseModel
 
@@ -58,11 +60,6 @@ class TestOllama(unittest.TestCase):
               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__':