⚡️ Speed up method TokenAuthClientProvider.authenticate by 406%

[codeflash-ai]

📄 406% (4.06x) speedup for TokenAuthClientProvider.authenticate in chromadb/auth/token_authn/__init__.py

⏱️ Runtime : 272 microseconds → 53.9 microseconds (best of 133 runs)

📝 Explanation and details

The optimization moves all expensive computations from the frequently-called authenticate() method to the constructor, achieving a 405% speedup by pre-computing the authentication headers once instead of on every call.

Key changes:

• Pre-computed authentication headers: The token retrieval (get_secret_value()), Bearer prefix formatting, and dictionary construction now happen once in __init__ and are stored in self._auth_header
• Simplified authenticate method: Now just returns the pre-computed self._auth_header instead of rebuilding it each time

Why this is faster:

• Eliminates repeated calls to self._token.get_secret_value() (20.6% of original runtime)
• Removes repeated string formatting for Bearer tokens (8.6% of original runtime)
• Avoids recreating the dictionary and SecretStr object on each call (58.2% of original runtime combined)
• Reduces total execution time from 1.147ms to 0.085ms per call

Test case performance:
The optimization shows consistent 385-566% speedups across all test scenarios, with particularly strong gains for:

• Large-scale tests with many tokens (398-415% faster)
• Long tokens up to 1000 characters (492-517% faster)
• All header types (Authorization and X-Chroma-Token both benefit equally)

This is a classic "compute once, use many times" optimization that's especially beneficial for authentication providers that may be called repeatedly during a session.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 784 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import string
# Enum for token header options
from enum import Enum
from typing import Dict

# imports
import pytest
from chromadb.auth.token_authn.__init__ import TokenAuthClientProvider
from pydantic import SecretStr

# --- Minimal stubs and enums to support the code under test ---


class TokenTransportHeader(Enum):
    AUTHORIZATION = "Authorization"
    X_CHROMA_TOKEN = "X-Chroma-Token"

# Minimal ClientAuthHeaders type
ClientAuthHeaders = Dict[str, SecretStr]

# Minimal System and Settings stubs
class Settings:
    def __init__(
        self,
        chroma_client_auth_credentials=None,
        chroma_auth_token_transport_header=None
    ):
        self.chroma_client_auth_credentials = chroma_client_auth_credentials
        self.chroma_auth_token_transport_header = chroma_auth_token_transport_header

    def require(self, key):
        val = getattr(self, key)
        if val is None:
            raise ValueError(f"Missing required config value '{key}'")
        return val

class System:
    def __init__(self, settings):
        self.settings = settings

# --- Code under test: authenticate() via TokenAuthClientProvider ---

valid_token_chars = set(string.digits + string.ascii_letters + string.punctuation)

def _check_token(token: str) -> None:
    token_str = str(token)
    if not all(c in valid_token_chars for c in token_str):
        raise ValueError(
            "Invalid token. Must contain only ASCII letters, digits, and punctuation."
        )

allowed_token_headers = [
    TokenTransportHeader.AUTHORIZATION.value,
    TokenTransportHeader.X_CHROMA_TOKEN.value,
]

def _check_allowed_token_headers(token_header: str) -> None:
    if token_header not in allowed_token_headers:
        raise ValueError(
            f"Invalid token transport header: {token_header}. "
            f"Must be one of {allowed_token_headers}"
        )

# Base class stub
class ClientAuthProvider:
    def __init__(self, system):
        pass
from chromadb.auth.token_authn.__init__ import TokenAuthClientProvider

# --- Pytest test suite ---

# ========== BASIC TEST CASES ==========

def test_authenticate_authorization_header_default():
    # Test default behavior: valid ASCII token, default header
    token = "ValidToken123!@#"
    settings = Settings(chroma_client_auth_credentials=token)
    system = System(settings)
    provider = TokenAuthClientProvider(system)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 2.09μs -> 384ns (445% faster)

def test_authenticate_x_chroma_token_header():
    # Test use of X-Chroma-Token header
    token = "Token456$%^"
    settings = Settings(
        chroma_client_auth_credentials=token,
        chroma_auth_token_transport_header="X-Chroma-Token"
    )
    system = System(settings)
    provider = TokenAuthClientProvider(system)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.64μs -> 297ns (452% faster)

def test_authenticate_authorization_header_specified():
    # Test explicit Authorization header
    token = "AnotherToken789"
    settings = Settings(
        chroma_client_auth_credentials=token,
        chroma_auth_token_transport_header="Authorization"
    )
    system = System(settings)
    provider = TokenAuthClientProvider(system)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.79μs -> 293ns (512% faster)

# ========== EDGE TEST CASES ==========

def test_authenticate_token_with_all_valid_ascii():
    # Token with all allowed characters
    token = string.ascii_letters + string.digits + string.punctuation
    settings = Settings(chroma_client_auth_credentials=token)
    system = System(settings)
    provider = TokenAuthClientProvider(system)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.64μs -> 263ns (524% faster)

def test_authenticate_token_with_invalid_unicode():
    # Token with invalid unicode character (e.g., emoji)
    token = "Valid123🙂"
    settings = Settings(chroma_client_auth_credentials=token)
    system = System(settings)
    with pytest.raises(ValueError, match="Invalid token"):
        TokenAuthClientProvider(system)


def test_authenticate_token_with_tab_newline():
    # Token with tab/newline (not in valid_token_chars, should fail)
    token = "TokenWithTab\t"
    settings = Settings(chroma_client_auth_credentials=token)
    system = System(settings)
    with pytest.raises(ValueError, match="Invalid token"):
        TokenAuthClientProvider(system)
    token2 = "TokenWithNewline\n"
    settings2 = Settings(chroma_client_auth_credentials=token2)
    system2 = System(settings2)
    with pytest.raises(ValueError, match="Invalid token"):
        TokenAuthClientProvider(system2)

def test_authenticate_missing_token():
    # Token is None (missing required config)
    settings = Settings(chroma_client_auth_credentials=None)
    system = System(settings)
    with pytest.raises(ValueError, match="Missing required config value 'chroma_client_auth_credentials'"):
        TokenAuthClientProvider(system)

def test_authenticate_empty_token():
    # Empty string token is technically valid (all chars are valid, even if 0)
    token = ""
    settings = Settings(chroma_client_auth_credentials=token)
    system = System(settings)
    provider = TokenAuthClientProvider(system)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 2.01μs -> 414ns (385% faster)

def test_authenticate_invalid_token_header():
    # Invalid token header specified
    token = "ValidToken123"
    settings = Settings(
        chroma_client_auth_credentials=token,
        chroma_auth_token_transport_header="X-Invalid-Header"
    )
    system = System(settings)
    with pytest.raises(ValueError, match="Invalid token transport header: X-Invalid-Header"):
        TokenAuthClientProvider(system)

def test_authenticate_token_header_case_sensitivity():
    # Header is case sensitive; "authorization" is not valid
    token = "ValidToken123"
    settings = Settings(
        chroma_client_auth_credentials=token,
        chroma_auth_token_transport_header="authorization"
    )
    system = System(settings)
    with pytest.raises(ValueError, match="Invalid token transport header: authorization"):
        TokenAuthClientProvider(system)

def test_authenticate_token_with_only_punctuation():
    # Token of only punctuation chars
    token = string.punctuation
    settings = Settings(chroma_client_auth_credentials=token)
    system = System(settings)
    provider = TokenAuthClientProvider(system)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.88μs -> 329ns (472% faster)

def test_authenticate_token_with_only_digits():
    # Token of only digits
    token = "0123456789"
    settings = Settings(chroma_client_auth_credentials=token)
    system = System(settings)
    provider = TokenAuthClientProvider(system)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.80μs -> 332ns (441% faster)

# ========== LARGE SCALE TEST CASES ==========

def test_authenticate_long_token():
    # Token at the upper reasonable length (e.g., 1000 chars)
    token = (string.ascii_letters + string.digits + string.punctuation) * 7
    token = token[:1000]  # ensure exactly 1000 chars
    settings = Settings(chroma_client_auth_credentials=token)
    system = System(settings)
    provider = TokenAuthClientProvider(system)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.83μs -> 309ns (492% faster)

def test_authenticate_many_instances():
    # Create many providers in a loop to check for leaks or state issues
    tokens = [f"Token{i:03d}!" for i in range(100)]
    for token in tokens:
        settings = Settings(chroma_client_auth_credentials=token)
        system = System(settings)
        provider = TokenAuthClientProvider(system)
        codeflash_output = provider.authenticate(); headers = codeflash_output # 66.0μs -> 13.3μs (398% faster)

def test_authenticate_many_headers():
    # Test both headers in a large loop
    tokens = [f"Token{i:03d}!" for i in range(50)]
    headers_list = ["Authorization", "X-Chroma-Token"]
    for i, token in enumerate(tokens):
        header = headers_list[i % 2]
        settings = Settings(
            chroma_client_auth_credentials=token,
            chroma_auth_token_transport_header=header
        )
        system = System(settings)
        provider = TokenAuthClientProvider(system)
        codeflash_output = provider.authenticate(); headers = codeflash_output # 34.1μs -> 6.84μs (399% faster)
        if header == "Authorization":
            pass
        else:
            pass

def test_authenticate_large_token_set_unique():
    # Use a set of unique tokens to check for cross-instance contamination
    tokens = [f"UniqueToken_{i}_{string.ascii_letters[i%52]}" for i in range(100)]
    providers = []
    for token in tokens:
        settings = Settings(chroma_client_auth_credentials=token)
        system = System(settings)
        provider = TokenAuthClientProvider(system)
        providers.append(provider)
    # Now authenticate and check all tokens are correct
    for i, provider in enumerate(providers):
        codeflash_output = provider.authenticate(); headers = codeflash_output # 63.1μs -> 13.0μs (385% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import string
# Enum for token transport headers
from enum import Enum
from typing import Dict

# imports
import pytest
from chromadb.auth.token_authn.__init__ import TokenAuthClientProvider

# --- Minimal stubs and implementations to support the test suite ---


class TokenTransportHeader(Enum):
    AUTHORIZATION = "Authorization"
    X_CHROMA_TOKEN = "X-Chroma-Token"

# SecretStr stub (from pydantic)
class SecretStr:
    def __init__(self, val: str):
        self._val = val
    def get_secret_value(self):
        return self._val
    def __eq__(self, other):
        if isinstance(other, SecretStr):
            return self._val == other._val
        if isinstance(other, str):
            return self._val == other
        return False
    def __repr__(self):
        return f"SecretStr({self._val})"

# Minimal System and Settings stubs
class Settings:
    def __init__(
        self,
        chroma_client_auth_credentials=None,
        chroma_auth_token_transport_header=None,
    ):
        self.chroma_client_auth_credentials = chroma_client_auth_credentials
        self.chroma_auth_token_transport_header = chroma_auth_token_transport_header

    def require(self, key):
        val = getattr(self, key)
        if val is None:
            raise ValueError(f"Missing required config value '{key}'")
        return val

class System:
    def __init__(self, settings):
        self.settings = settings

# ClientAuthProvider stub
class ClientAuthProvider:
    def __init__(self, system):
        pass

# --- Code under test: TokenAuthClientProvider and helpers ---

valid_token_chars = set(string.digits + string.ascii_letters + string.punctuation)
allowed_token_headers = [
    TokenTransportHeader.AUTHORIZATION.value,
    TokenTransportHeader.X_CHROMA_TOKEN.value,
]

def _check_token(token: str) -> None:
    token_str = str(token)
    if not all(c in valid_token_chars for c in token_str):
        raise ValueError(
            "Invalid token. Must contain only ASCII letters, digits, and punctuation."
        )

def _check_allowed_token_headers(token_header: str) -> None:
    if token_header not in allowed_token_headers:
        raise ValueError(
            f"Invalid token transport header: {token_header}. "
            f"Must be one of {allowed_token_headers}"
        )
from chromadb.auth.token_authn.__init__ import TokenAuthClientProvider

# --- Unit tests for TokenAuthClientProvider.authenticate ---

# -------------------
# 1. Basic Test Cases
# -------------------

def test_authenticate_default_header_bearer():
    """Test with a valid ASCII token and default header (Authorization)"""
    token = "ValidToken123!@#"
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    provider = TokenAuthClientProvider(sys)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.57μs -> 238ns (561% faster)

def test_authenticate_x_chroma_token_header():
    """Test with a valid token and X-Chroma-Token header"""
    token = "AnotherValidToken_456"
    s = Settings(
        chroma_client_auth_credentials=token,
        chroma_auth_token_transport_header="X-Chroma-Token",
    )
    sys = System(s)
    provider = TokenAuthClientProvider(sys)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.51μs -> 248ns (508% faster)

def test_authenticate_authorization_header_explicit():
    """Test with explicit Authorization header and valid token"""
    token = "TokenWithExplicitHeader"
    s = Settings(
        chroma_client_auth_credentials=token,
        chroma_auth_token_transport_header="Authorization",
    )
    sys = System(s)
    provider = TokenAuthClientProvider(sys)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.66μs -> 237ns (602% faster)

# -------------------
# 2. Edge Test Cases
# -------------------

def test_token_with_all_valid_ascii_characters():
    """Test a token containing all ASCII letters, digits, and punctuation"""
    token = string.ascii_letters + string.digits + string.punctuation
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    provider = TokenAuthClientProvider(sys)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.60μs -> 271ns (491% faster)

def test_token_with_invalid_unicode_character():
    """Test a token containing a non-ASCII character (should raise ValueError)"""
    token = "Valid123" + "💥"  # includes emoji
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    with pytest.raises(ValueError, match="Invalid token"):
        TokenAuthClientProvider(sys)

def test_token_with_whitespace():
    """Test a token containing whitespace (should be valid, as whitespace is in string.punctuation)"""
    # Actually, whitespace is not in string.punctuation; let's check with a space
    token = "ValidToken WithSpace"
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    with pytest.raises(ValueError, match="Invalid token"):
        TokenAuthClientProvider(sys)

def test_token_with_control_character():
    """Test a token containing a control character (e.g., newline)"""
    token = "ValidToken\n"
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    with pytest.raises(ValueError, match="Invalid token"):
        TokenAuthClientProvider(sys)

def test_missing_token_raises():
    """Test that missing chroma_client_auth_credentials raises ValueError"""
    s = Settings()
    sys = System(s)
    with pytest.raises(ValueError, match="Missing required config value 'chroma_client_auth_credentials'"):
        TokenAuthClientProvider(sys)

def test_invalid_header_raises():
    """Test that an invalid header raises ValueError"""
    token = "ValidToken123"
    s = Settings(
        chroma_client_auth_credentials=token,
        chroma_auth_token_transport_header="Invalid-Header",
    )
    sys = System(s)
    with pytest.raises(ValueError, match="Invalid token transport header"):
        TokenAuthClientProvider(sys)

def test_case_sensitivity_of_header():
    """Test that header is case sensitive (should raise if not exact match)"""
    token = "ValidToken"
    s = Settings(
        chroma_client_auth_credentials=token,
        chroma_auth_token_transport_header="authorization",  # lower case, should fail
    )
    sys = System(s)
    with pytest.raises(ValueError, match="Invalid token transport header"):
        TokenAuthClientProvider(sys)

def test_token_empty_string():
    """Test that an empty string token is allowed (since all chars in empty string are valid)"""
    token = ""
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    provider = TokenAuthClientProvider(sys)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.92μs -> 353ns (445% faster)

def test_token_with_only_punctuation():
    """Test a token with only punctuation characters"""
    token = string.punctuation
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    provider = TokenAuthClientProvider(sys)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.79μs -> 308ns (482% faster)

# -------------------------
# 3. Large Scale Test Cases
# -------------------------

def test_long_token_approaching_1000_chars():
    """Test with a long token (1000 chars) of valid characters"""
    token = (string.ascii_letters + string.digits + string.punctuation) * 7
    token = token[:1000]  # ensure exactly 1000 chars
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    provider = TokenAuthClientProvider(sys)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.82μs -> 295ns (517% faster)

def test_large_batch_of_tokens():
    """Test authenticating with many different valid tokens in a batch (performance + correctness)"""
    base = string.ascii_letters + string.digits
    for i in range(100):  # 100 different tokens, each unique
        token = f"Token{i}_" + base[i % len(base)] * (10 + i)
        s = Settings(chroma_client_auth_credentials=token)
        sys = System(s)
        provider = TokenAuthClientProvider(sys)
        codeflash_output = provider.authenticate(); headers = codeflash_output # 66.5μs -> 13.1μs (408% faster)

def test_many_headers_in_parallel():
    """Test that many providers with different headers work in parallel"""
    tokens = [f"Token_{i}_!@#" for i in range(20)]
    headers = ["Authorization", "X-Chroma-Token"]
    for i, token in enumerate(tokens):
        header = headers[i % 2]
        s = Settings(
            chroma_client_auth_credentials=token,
            chroma_auth_token_transport_header=header,
        )
        sys = System(s)
        provider = TokenAuthClientProvider(sys)
        codeflash_output = provider.authenticate(); out = codeflash_output # 14.5μs -> 2.82μs (415% faster)
        if header == "Authorization":
            pass
        else:
            pass

def test_token_with_all_punctuation_many_times():
    """Test with a token that is all punctuation, repeated to near 1000 chars"""
    token = (string.punctuation * (1000 // len(string.punctuation) + 1))[:1000]
    s = Settings(chroma_client_auth_credentials=token)
    sys = System(s)
    provider = TokenAuthClientProvider(sys)
    codeflash_output = provider.authenticate(); headers = codeflash_output # 1.70μs -> 255ns (566% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-TokenAuthClientProvider.authenticate-mh7m3m9g and push.