From 1a60fa16674ee5d70a359018dbfd545f6689cbfe Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 5 Nov 2024 21:50:00 -0800
Subject: [PATCH] Test to exercise gpt-4o-audio-preview, closes #608

---
 tests/test_cli_openai_models.py | 88 +++++++++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)

diff --git a/tests/test_cli_openai_models.py b/tests/test_cli_openai_models.py
index 56acd11..f341e38 100644
--- a/tests/test_cli_openai_models.py
+++ b/tests/test_cli_openai_models.py
@@ -56,3 +56,91 @@ def test_openai_options_min_max():
         result2 = runner.invoke(cli, ["-m", "chatgpt", "-o", option, "10"])
         assert result2.exit_code == 1
         assert f"less than or equal to {max_val}" in result2.output
+
+
+@pytest.mark.parametrize("model", ("gpt-4o-mini", "gpt-4o-audio-preview"))
+@pytest.mark.parametrize("filetype", ("mp3", "wav"))
+def test_only_gpt4_audio_preview_allows_mp3_or_wav(httpx_mock, model, filetype):
+    httpx_mock.add_response(
+        method="HEAD",
+        url=f"https://www.example.com/example.{filetype}",
+        content=b"binary-data",
+        headers={"Content-Type": "audio/mpeg" if filetype == "mp3" else "audio/wave"},
+    )
+    # Another mock for the correct model
+    if model == "gpt-4o-audio-preview":
+        httpx_mock.add_response(
+            method="POST",
+            # chat completion request
+            url="https://api.openai.com/v1/chat/completions",
+            json={
+                "id": "chatcmpl-AQT9a30kxEaM1bqxRPepQsPlCyGJh",
+                "object": "chat.completion",
+                "created": 1730871958,
+                "model": "gpt-4o-audio-preview-2024-10-01",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": "Why did the pelican get kicked out of the restaurant?\n\nBecause he had a big bill and no way to pay it!",
+                            "refusal": None,
+                        },
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 55,
+                    "completion_tokens": 25,
+                    "total_tokens": 80,
+                    "prompt_tokens_details": {
+                        "cached_tokens": 0,
+                        "audio_tokens": 44,
+                        "text_tokens": 11,
+                        "image_tokens": 0,
+                    },
+                    "completion_tokens_details": {
+                        "reasoning_tokens": 0,
+                        "audio_tokens": 0,
+                        "text_tokens": 25,
+                        "accepted_prediction_tokens": 0,
+                        "rejected_prediction_tokens": 0,
+                    },
+                },
+                "system_fingerprint": "fp_49254d0e9b",
+            },
+            headers={"Content-Type": "application/json"},
+        )
+        httpx_mock.add_response(
+            method="GET",
+            url=f"https://www.example.com/example.{filetype}",
+            content=b"binary-data",
+            headers={
+                "Content-Type": "audio/mpeg" if filetype == "mp3" else "audio/wave"
+            },
+        )
+    runner = CliRunner()
+    result = runner.invoke(
+        cli,
+        [
+            "-m",
+            model,
+            "-a",
+            f"https://www.example.com/example.{filetype}",
+            "--no-stream",
+            "--key",
+            "x",
+        ],
+    )
+    if model == "gpt-4o-audio-preview":
+        assert result.exit_code == 0
+        assert result.output == (
+            "Why did the pelican get kicked out of the restaurant?\n\n"
+            "Because he had a big bill and no way to pay it!\n"
+        )
+    else:
+        assert result.exit_code == 1
+        long = "audio/mpeg" if filetype == "mp3" else "audio/wave"
+        assert (
+            f"This model does not support attachments of type '{long}'" in result.output
+        )