From a9e2ef62125c462cc62d824f3d90bbc1d0366dfe Mon Sep 17 00:00:00 2001
From: Lucca Zenobio <lucca.zenobio@ifood.com.br>
Date: Mon, 29 Apr 2024 10:05:30 -0300
Subject: [PATCH] test

---
 litellm/tests/test_bedrock_completion.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py
index ca2ffea5f5..2aab8a3b44 100644
--- a/litellm/tests/test_bedrock_completion.py
+++ b/litellm/tests/test_bedrock_completion.py
@@ -207,6 +207,25 @@ def test_completion_bedrock_claude_sts_client_auth():
 # test_completion_bedrock_claude_sts_client_auth()
 
 
+def test_bedrock_extra_headers():
+    try:
+        litellm.set_verbose = True
+        response: ModelResponse = completion(
+            model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+            messages=messages,
+            max_tokens=10,
+            temperature=0.78,
+            extra_headers={"x-key": "x_key_value"}
+        )
+        # Add any assertions here to check the response
+        assert len(response.choices) > 0
+        assert len(response.choices[0].message.content) > 0
+    except RateLimitError:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+
 def test_bedrock_claude_3():
     try:
         litellm.set_verbose = True