From 2a5966eb10415cdbe1b09d270e2a0a1e5a272c59 Mon Sep 17 00:00:00 2001
From: Stan Zubarev <szubarev@gmail.com>
Date: Thu, 31 Oct 2024 20:50:42 -0400
Subject: [PATCH] fix tests

---
 tests/test_generate_data.py | 3 ++-
 tests/test_generate_text.py | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/test_generate_data.py b/tests/test_generate_data.py
index bea35c8..b622cad 100644
--- a/tests/test_generate_data.py
+++ b/tests/test_generate_data.py
@@ -1,6 +1,6 @@
 import pytest
 
-from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama
+from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama, Amazon
 from pydantic import BaseModel
 
 
@@ -16,6 +16,7 @@ class ResponseModel(BaseModel):
         OpenAI,
         Groq,
         Ollama,
+        Amazon
     ],
 )
 def test_generate_data(provider_cls):
diff --git a/tests/test_generate_text.py b/tests/test_generate_text.py
index 55a9a13..0611b1d 100644
--- a/tests/test_generate_text.py
+++ b/tests/test_generate_text.py
@@ -1,6 +1,6 @@
 import pytest
 
-from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama
+from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama, Amazon
 
 
 @pytest.mark.parametrize(
@@ -11,6 +11,7 @@ from simplemind.providers import Anthropic, Gemini, OpenAI, Groq, Ollama
         OpenAI,
         Groq,
         Ollama,
+        Amazon,
     ],
 )
 def test_generate_text(provider_cls):