Skip to content

Commit

Permalink
fix: tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Josephasafg committed Sep 19, 2024
1 parent f00726c commit c43c0e9
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 10 deletions.
3 changes: 2 additions & 1 deletion tests/unittests/test_ai21_env_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
from ai21 import AI21Client

_FAKE_API_KEY = "fake-key"
os.environ["AI21_API_KEY"] = _FAKE_API_KEY


@contextmanager
Expand All @@ -15,6 +14,7 @@ def set_env_var(key: str, value: str):


def test_env_config__when_set_via_init_and_env__should_be_taken_from_init():
os.environ["AI21_API_KEY"] = _FAKE_API_KEY
client = AI21Client()
assert client._api_key == _FAKE_API_KEY

Expand All @@ -25,6 +25,7 @@ def test_env_config__when_set_via_init_and_env__should_be_taken_from_init():


def test_env_config__when_set_twice__should_be_updated():
os.environ["AI21_API_KEY"] = _FAKE_API_KEY
client = AI21Client()

assert client._api_key == _FAKE_API_KEY
Expand Down
17 changes: 8 additions & 9 deletions tests/unittests/test_http_client.py
Original file line number Diff line number Diff line change
@@ -1,21 +1,20 @@
import pytest
from unittest.mock import Mock
from urllib.request import Request

import httpx
import pytest

from ai21.errors import ServiceUnavailable, Unauthorized
from ai21.http_client.async_http_client import AsyncAI21HTTPClient
from ai21.http_client.base_http_client import RETRY_ERROR_CODES
from ai21.http_client.http_client import AI21HTTPClient
from ai21.http_client.async_http_client import AsyncAI21HTTPClient

_METHOD = "GET"
_URL = "http://test_url"
_API_KEY = "fake-key"


def test__execute_http_request__when_retry_error_code_once__should_retry_and_succeed(mock_httpx_client: Mock) -> None:
request = Request(method=_METHOD, url=_URL)
request = httpx.Request(method=_METHOD, url=_URL)
retries = 3
mock_httpx_client.send.side_effect = [
httpx.Response(status_code=429, request=request),
Expand All @@ -28,7 +27,7 @@ def test__execute_http_request__when_retry_error_code_once__should_retry_and_suc


def test__execute_http_request__when_retry_error__should_retry_and_stop(mock_httpx_client: Mock) -> None:
request = Request(method=_METHOD, url=_URL)
request = httpx.Request(method=_METHOD, url=_URL)
retries = len(RETRY_ERROR_CODES)

mock_httpx_client.send.side_effect = [
Expand All @@ -44,7 +43,7 @@ def test__execute_http_request__when_retry_error__should_retry_and_stop(mock_htt

def test__execute_http_request__when_streaming__should_handle_non_200_response_code(mock_httpx_client: Mock) -> None:
error_details = "test_error"
request = Request(method=_METHOD, url=_URL)
request = httpx.Request(method=_METHOD, url=_URL)
response = httpx.Response(status_code=401, request=request, text=error_details)
mock_httpx_client.send.return_value = response

Expand All @@ -57,7 +56,7 @@ def test__execute_http_request__when_streaming__should_handle_non_200_response_c
async def test__execute_async_http_request__when_retry_error_code_once__should_retry_and_succeed(
mock_httpx_async_client: Mock,
) -> None:
request = Request(method=_METHOD, url=_URL)
request = httpx.Request(method=_METHOD, url=_URL)
retries = 3
mock_httpx_async_client.send.side_effect = [
httpx.Response(status_code=429, request=request),
Expand All @@ -73,7 +72,7 @@ async def test__execute_async_http_request__when_retry_error_code_once__should_r
async def test__execute_async_http_request__when_retry_error__should_retry_and_stop(
mock_httpx_async_client: Mock,
) -> None:
request = Request(method=_METHOD, url=_URL)
request = httpx.Request(method=_METHOD, url=_URL)
retries = len(RETRY_ERROR_CODES)

mock_httpx_async_client.send.side_effect = [
Expand All @@ -92,7 +91,7 @@ async def test__execute_async_http_request__when_streaming__should_handle_non_20
mock_httpx_async_client: Mock,
) -> None:
error_details = "test_error"
request = Request(method=_METHOD, url=_URL)
request = httpx.Request(method=_METHOD, url=_URL)
response = httpx.Response(status_code=401, request=request, text=error_details)
mock_httpx_async_client.send.return_value = response

Expand Down

0 comments on commit c43c0e9

Please sign in to comment.