Repository: openai/openai-python
File: tests/test_client.py (Python)

Code before:
        assert response.json() == {"status": "ok"}

    @pytest.mark.respx(base_url=base_url)
    def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
        # Test that follow_redirects=False prevents following redirects
        respx_mock.post("/redirect").mock(
            return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
        )

        with pytest.raises(APIStatusError) as exc_info:
            self.client.post(
                "/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response
            )

        assert exc_info.value.response.status_code == 302
        assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"

    def test_refresh_auth_headers_token(self) -> None:
        client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token")
        client.refresh_auth_headers()
        assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"

    def test_refresh_auth_headers_key(self) -> None:
        client = OpenAI(base_url=base_url, api_key="test_api_key")
        client.refresh_auth_headers()
        assert client.auth_headers.get("Authorization") == "Bearer test_api_key"

    @pytest.mark.respx()
    def test_bearer_token_refresh(self, respx_mock: MockRouter) -> None:
        respx_mock.post(base_url + "/chat/completions").mock(
            side_effect=[
                httpx.Response(500, json={"error": "server error"}),
                httpx.Response(200, json={"foo": "bar"}),
            ]
        )

        counter = 0

        def token_provider() -> str:
            nonlocal counter

            counter += 1

            if counter == 1:
                return "first"

            return "second"

        client = OpenAI(base_url=base_url, api_key=token_provider)
        client.chat.completions.create(messages=[], model="gpt-4")


Review comment (security):
nit: I don't think these tests are super valuable as we're testing the runtime behaviour below.

a more useful test here would instead be what happens if you access `client.auth_headers` _before_ a request is made

Code after:
        assert response.json() == {"status": "ok"}

    @pytest.mark.respx(base_url=base_url)
    def test_follow_redirects_disabled(self, respx_mock: MockRouter) -> None:
        # Test that follow_redirects=False prevents following redirects
        respx_mock.post("/redirect").mock(
            return_value=httpx.Response(302, headers={"Location": f"{base_url}/redirected"})
        )

        with pytest.raises(APIStatusError) as exc_info:
            self.client.post(
                "/redirect", body={"key": "value"}, options={"follow_redirects": False}, cast_to=httpx.Response
            )

        assert exc_info.value.response.status_code == 302
        assert exc_info.value.response.headers["Location"] == f"{base_url}/redirected"

    def test_api_key_before_after_refresh_provider(self) -> None:
        client = OpenAI(base_url=base_url, api_key=lambda: "test_bearer_token")

        assert client.api_key == ""
        assert "Authorization" not in client.auth_headers

        client._refresh_api_key()

        assert client.api_key == "test_bearer_token"
        assert client.auth_headers.get("Authorization") == "Bearer test_bearer_token"

    def test_api_key_before_after_refresh_str(self) -> None:
        client = OpenAI(base_url=base_url, api_key="test_api_key")

        assert client.auth_headers.get("Authorization") == "Bearer test_api_key"
        client._refresh_api_key()

        assert client.auth_headers.get("Authorization") == "Bearer test_api_key"

    @pytest.mark.respx()
    def test_api_key_refresh_on_retry(self, respx_mock: MockRouter) -> None:
        respx_mock.post(base_url + "/chat/completions").mock(
            side_effect=[
                httpx.Response(500, json={"error": "server error"}),
                httpx.Response(200, json={"foo": "bar"}),
            ]
        )

        counter = 0

        def token_provider() -> str:
            nonlocal counter

            counter += 1