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

Code before:

        counter += 1

        if counter == 1:
            return "first"

        return "second"

    client = AsyncAzureOpenAI(
        api_version="2024-02-01",
        azure_ad_token_provider=token_provider,
        azure_endpoint="https://example-resource.azure.openai.com",
    )

    await client.chat.completions.create(messages=[], model="gpt-4")

    calls = cast("list[MockRequestCall]", respx_mock.calls)

    assert len(calls) == 2

    assert calls[0].request.headers.get("Authorization") == "Bearer first"
    assert calls[1].request.headers.get("Authorization") == "Bearer second"


@pytest.mark.respx()
def test_azure_api_key_redacted(respx_mock: MockRouter, _logger_with_filter: logging.Logger, caplog: pytest.LogCaptureFixture) -> None:
    respx_mock.post(
        "https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-06-01"
    ).mock(
        return_value=httpx.Response(200, json={"model": "gpt-4"})
    )

    client = AzureOpenAI(
        api_version="2024-06-01",
        api_key="example_api_key",
        azure_endpoint="https://example-resource.azure.openai.com",
    )

    with caplog.at_level(logging.DEBUG):
        client.chat.completions.create(messages=[], model="gpt-4")

    for record in caplog.records:
        if is_dict(record.args) and "headers" in record.args and is_dict(record.args["headers"]):
            assert record.args["headers"]["api-key"] == "<redacted>"


@pytest.mark.respx()
def test_azure_bearer_token_redacted(respx_mock: MockRouter, _logger_with_filter: logging.Logger, caplog: pytest.LogCaptureFixture) -> None:
    respx_mock.post(
        "https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-06-01"
    ).mock(

Review comment (suggestion):
note: I *think* you could've avoided the `_logger_with_filter` stuff by putting all these tests in a class and then doing
```py
@pytest.fixture(autouse=True)
def _logger_with_filter() -> logging.Logger: ...
```
but totally not worth blocking on that change

Code after:
    client = AsyncAzureOpenAI(
        api_version="2024-02-01",
        azure_ad_token_provider=token_provider,
        azure_endpoint="https://example-resource.azure.openai.com",
    )

    await client.chat.completions.create(messages=[], model="gpt-4")

    calls = cast("list[MockRequestCall]", respx_mock.calls)

    assert len(calls) == 2

    assert calls[0].request.headers.get("Authorization") == "Bearer first"
    assert calls[1].request.headers.get("Authorization") == "Bearer second"


class TestAzureLogging:

    @pytest.fixture(autouse=True)
    def logger_with_filter(self) -> logging.Logger:
        logger = logging.getLogger("openai")
        logger.setLevel(logging.DEBUG)
        logger.addFilter(SensitiveHeadersFilter())
        return logger

    @pytest.mark.respx()
    def test_azure_api_key_redacted(self, respx_mock: MockRouter, caplog: pytest.LogCaptureFixture) -> None:
        respx_mock.post(
            "https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-06-01"
        ).mock(
            return_value=httpx.Response(200, json={"model": "gpt-4"})
        )

        client = AzureOpenAI(
            api_version="2024-06-01",
            api_key="example_api_key",
            azure_endpoint="https://example-resource.azure.openai.com",
        )

        with caplog.at_level(logging.DEBUG):
            client.chat.completions.create(messages=[], model="gpt-4")

        for record in caplog.records:
            if is_dict(record.args) and record.args.get("headers") and is_dict(record.args["headers"]):
                assert record.args["headers"]["api-key"] == "<redacted>"


    @pytest.mark.respx()
    def test_azure_bearer_token_redacted(self, respx_mock: MockRouter, caplog: pytest.LogCaptureFixture) -> None:
        respx_mock.post(
            "https://example-resource.azure.openai.com/openai/deployments/gpt-4/chat/completions?api-version=2024-06-01"