From 58ed6e77de83f3f3e7f74947c8e4195173553870 Mon Sep 17 00:00:00 2001 From: Tim Xia Date: Fri, 1 Mar 2024 23:27:02 -0500 Subject: [PATCH] add assertion for test --- litellm/tests/test_bedrock_completion.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index e9ad6ac1c34..6f322207428 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -11,7 +11,7 @@ sys.path.insert( ) # Adds the parent directory to the system path import pytest import litellm -from litellm import embedding, completion, completion_cost, Timeout +from litellm import embedding, completion, completion_cost, Timeout, ModelResponse from litellm import RateLimitError # litellm.num_retries = 3 @@ -270,14 +270,15 @@ def test_completion_bedrock_mistral_completion_auth(): # os.environ.pop("AWS_SECRET_ACCESS_KEY", None) # os.environ.pop("AWS_REGION_NAME", None) try: - response = completion( + response:ModelResponse = completion( model="bedrock/mistral.mistral-7b-instruct-v0:2", messages=messages, max_tokens=10, temperature=0.1, ) # Add any assertions here to check the response - print(response) + assert len(response.choices) > 0 + assert len(response.choices[0].message.content) > 0 # os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id # os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key