From 049915c14bcea13666707a82c231db8a7eb46d5f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 22 Jan 2025 18:52:11 -0800
Subject: [PATCH] test: mock fireworks ai test - unstable api

---
 .../test_fireworks_ai_translation.py          | 51 +++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/tests/llm_translation/test_fireworks_ai_translation.py b/tests/llm_translation/test_fireworks_ai_translation.py
index 5e1a1c05120..9e78270c922 100644
--- a/tests/llm_translation/test_fireworks_ai_translation.py
+++ b/tests/llm_translation/test_fireworks_ai_translation.py
@@ -93,6 +93,57 @@ class TestFireworksAIChatCompletion(BaseLLMChatTest):
         """
         pass
 
+    @pytest.mark.parametrize(
+        "response_format",
+        [
+            {"type": "json_object"},
+            {"type": "text"},
+        ],
+    )
+    @pytest.mark.flaky(retries=6, delay=1)
+    def test_json_response_format(self, response_format):
+        """
+        Test that the JSON response format is supported by the LLM API
+        """
+        from litellm.utils import supports_response_schema
+        from openai import OpenAI
+        from unittest.mock import patch
+
+        client = OpenAI()
+
+        base_completion_call_args = self.get_base_completion_call_args()
+        litellm.set_verbose = True
+
+        messages = [
+            {
+                "role": "system",
+                "content": "Your output should be a JSON object with no additional properties.  ",
+            },
+            {
+                "role": "user",
+                "content": "Respond with this in json. city=San Francisco, state=CA, weather=sunny, temp=60",
+            },
+        ]
+
+        with patch.object(
+            client.chat.completions.with_raw_response, "create"
+        ) as mock_post:
+            response = self.completion_function(
+                **base_completion_call_args,
+                messages=messages,
+                response_format=response_format,
+                client=client,
+            )
+
+            mock_post.assert_called_once()
+            if response_format["type"] == "json_object":
+                assert (
+                    mock_post.call_args.kwargs["response_format"]["type"]
+                    == "json_object"
+                )
+            else:
+                assert mock_post.call_args.kwargs["response_format"]["type"] == "text"
+
 
 class TestFireworksAIAudioTranscription(BaseLLMAudioTranscriptionTest):
     def get_base_audio_transcription_call_args(self) -> dict: