⚡️ Speed up function default_api_key by 59% - #5

from __future__ import annotations

import os

# imports
import pytest
from together.utils.api_helpers import default_api_key


# Minimal stubs for together.error and together.constants for testing
class AuthenticationError(Exception):
    pass

class constants:
    MISSING_API_KEY_MESSAGE = "API key is missing."

class error:
    AuthenticationError = AuthenticationError
from together.utils.api_helpers import default_api_key

# 1. Basic Test Cases

def test_returns_supplied_api_key(monkeypatch):
    """If api_key is supplied, it is returned (env var ignored)."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key("supplied_key"); result = codeflash_output # 291ns -> 335ns (13.1% slower)

def test_returns_env_api_key(monkeypatch):
    """If api_key is None, returns the env var."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key(); result = codeflash_output # 1.61μs -> 1.23μs (31.3% faster)

def test_returns_env_api_key_when_api_key_is_empty_string(monkeypatch):
    """If api_key is '', uses the env var."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key(""); result = codeflash_output # 1.53μs -> 1.06μs (43.7% faster)

def test_returns_env_api_key_when_api_key_is_falsey(monkeypatch):
    """If api_key is a falsey value (except None), uses the env var."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    for falsey in [None, "", 0, False]:
        # Only None and '' should trigger env var fallback
        if falsey is None or falsey == "":
            codeflash_output = default_api_key(falsey)
        else:
            # 0 and False are not valid API keys, but will be treated as not supplied, so fallback to env
            codeflash_output = default_api_key(falsey)


def test_api_key_priority_over_env(monkeypatch):
    """Supplied api_key takes priority over env var, even if both are set."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key("supplied_key"); result = codeflash_output # 395ns -> 410ns (3.66% slower)


def test_api_key_is_whitespace(monkeypatch):
    """Whitespace string as api_key is truthy and returned as is."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key("   "); result = codeflash_output # 404ns -> 377ns (7.16% faster)

def test_env_api_key_is_whitespace(monkeypatch):
    """Whitespace string as env var is truthy and returned as is."""
    monkeypatch.setenv("TOGETHER_API_KEY", "   ")
    codeflash_output = default_api_key(); result = codeflash_output # 1.66μs -> 1.17μs (41.5% faster)

def test_api_key_is_zero(monkeypatch):
    """api_key=0 is falsey, so falls back to env var."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key(0) # 1.48μs -> 1.09μs (35.9% faster)


def test_api_key_is_false(monkeypatch):
    """api_key=False is falsey, so falls back to env var."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key(False) # 1.75μs -> 1.17μs (49.6% faster)

def test_api_key_is_object(monkeypatch):
    """Non-string object as api_key is truthy and returned as is."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    obj = object()
    codeflash_output = default_api_key(obj) # 317ns -> 339ns (6.49% slower)

def test_api_key_is_list(monkeypatch):
    """Non-empty list as api_key is truthy and returned as is."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    key = ["a", "b"]
    codeflash_output = default_api_key(key) # 324ns -> 327ns (0.917% slower)

def test_api_key_is_empty_list(monkeypatch):
    """Empty list as api_key is falsey, so falls back to env var."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key([]) # 1.61μs -> 1.15μs (39.9% faster)

def test_env_api_key_is_numeric_string(monkeypatch):
    """Numeric string as env var is truthy and returned as is."""
    monkeypatch.setenv("TOGETHER_API_KEY", "123456")
    codeflash_output = default_api_key() # 1.57μs -> 1.15μs (36.3% faster)

# 3. Large Scale Test Cases

def test_env_api_key_large_string(monkeypatch):
    """Very large API key in env var is returned as is."""
    large_key = "A" * 1000
    monkeypatch.setenv("TOGETHER_API_KEY", large_key)
    codeflash_output = default_api_key() # 2.06μs -> 1.38μs (48.8% faster)

def test_supplied_api_key_large_string(monkeypatch):
    """Very large supplied API key is returned as is."""
    large_key = "B" * 1000
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key(large_key) # 302ns -> 299ns (1.00% faster)

def test_many_calls_with_different_env(monkeypatch):
    """Test function with many different env var values for robustness."""
    for i in range(100):
        key = f"env_key_{i}"
        monkeypatch.setenv("TOGETHER_API_KEY", key)
        codeflash_output = default_api_key() # 87.9μs -> 51.2μs (71.5% faster)
        # Supplied API key always takes priority
        codeflash_output = default_api_key(f"supplied_{i}")


def test_api_key_various_types(monkeypatch):
    """Test many types as api_key to ensure only falsey values fall back to env."""
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    test_cases = [
        ("", "env_key"),
        (None, "env_key"),
        ([], "env_key"),
        ({}, "env_key"),
        (0, "env_key"),
        (False, "env_key"),
        ("nonempty", "nonempty"),
        ([1,2], [1,2]),
        ({"a":1}, {"a":1}),
        (True, True),
        (1, 1),
    ]
    for supplied, expected in test_cases:
        codeflash_output = default_api_key(supplied) # 7.21μs -> 5.00μs (44.2% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from __future__ import annotations

import os

# imports
import pytest
import together
from together import error
from together.utils.api_helpers import default_api_key

# unit tests

# --- Basic Test Cases ---

def test_api_key_argument_priority(monkeypatch):
    # If api_key argument is provided, it should be returned regardless of env var
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key("argument_key"); result = codeflash_output # 289ns -> 298ns (3.02% slower)

def test_api_key_env_var(monkeypatch):
    # If api_key argument is None, should return env var if present
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key(); result = codeflash_output # 1.52μs -> 1.13μs (34.2% faster)

def test_api_key_argument_over_env_var(monkeypatch):
    # If both are present, argument takes priority
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key("argument_key"); result = codeflash_output # 272ns -> 304ns (10.5% slower)

def test_api_key_none_and_no_env(monkeypatch):
    # If neither argument nor env var, should raise AuthenticationError
    monkeypatch.delenv("TOGETHER_API_KEY", raising=False)
    with pytest.raises(error.AuthenticationError):
        default_api_key() # 6.27μs -> 5.95μs (5.41% faster)

# --- Edge Test Cases ---

def test_api_key_empty_string(monkeypatch):
    # If api_key argument is empty string, should treat as not provided and use env var
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key(""); result = codeflash_output # 1.58μs -> 1.12μs (41.2% faster)

def test_env_var_empty_string(monkeypatch):
    # If env var is empty string, should raise AuthenticationError
    monkeypatch.setenv("TOGETHER_API_KEY", "")
    with pytest.raises(error.AuthenticationError):
        default_api_key() # 4.77μs -> 4.68μs (2.01% faster)

def test_api_key_falsey_values(monkeypatch):
    # Falsey values other than None/empty string (e.g. 0, False) should fallback to env var
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    for falsey in [0, False]:
        codeflash_output = default_api_key(falsey); result = codeflash_output # 2.56μs -> 1.83μs (40.1% faster)

def test_api_key_none_and_env_var_none(monkeypatch):
    # If both argument and env var are None, should raise AuthenticationError
    monkeypatch.delenv("TOGETHER_API_KEY", raising=False)
    with pytest.raises(error.AuthenticationError):
        default_api_key(None) # 4.70μs -> 4.76μs (1.28% slower)

def test_api_key_argument_is_whitespace(monkeypatch):
    # Whitespace string as api_key should fallback to env var
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key("   "); result = codeflash_output # 325ns -> 344ns (5.52% slower)

def test_env_var_whitespace(monkeypatch):
    # Whitespace string as env var should be returned if argument is not provided
    monkeypatch.setenv("TOGETHER_API_KEY", "   ")
    codeflash_output = default_api_key(); result = codeflash_output # 1.54μs -> 1.08μs (42.2% faster)

def test_api_key_argument_is_integer(monkeypatch):
    # Integer argument should be returned as is (converted to string if needed)
    monkeypatch.delenv("TOGETHER_API_KEY", raising=False)
    codeflash_output = default_api_key(123456); result = codeflash_output # 320ns -> 315ns (1.59% faster)

def test_api_key_argument_is_bytes(monkeypatch):
    # Bytes argument should be returned as is
    monkeypatch.delenv("TOGETHER_API_KEY", raising=False)
    codeflash_output = default_api_key(b"keybytes"); result = codeflash_output # 306ns -> 327ns (6.42% slower)

# --- Large Scale Test Cases ---

def test_many_env_var_calls(monkeypatch):
    # Test repeated calls with env var set
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    for _ in range(1000):
        codeflash_output = default_api_key() # 837μs -> 474μs (76.4% faster)

def test_many_argument_calls(monkeypatch):
    # Test repeated calls with argument provided
    monkeypatch.delenv("TOGETHER_API_KEY", raising=False)
    for i in range(1000):
        key = f"key_{i}"
        codeflash_output = default_api_key(key) # 111μs -> 111μs (0.382% faster)

def test_large_env_var(monkeypatch):
    # Test with a large API key in env var
    large_key = "A" * 1000
    monkeypatch.setenv("TOGETHER_API_KEY", large_key)
    codeflash_output = default_api_key(); result = codeflash_output # 2.21μs -> 1.63μs (35.6% faster)

def test_large_argument(monkeypatch):
    # Test with a large API key as argument
    large_key = "B" * 1000
    monkeypatch.delenv("TOGETHER_API_KEY", raising=False)
    codeflash_output = default_api_key(large_key); result = codeflash_output # 283ns -> 292ns (3.08% slower)

def test_large_scale_argument_and_env(monkeypatch):
    # Test with both large argument and large env var, argument should be returned
    arg_key = "ARG" * 333 + "A"
    env_key = "ENV" * 333 + "E"
    monkeypatch.setenv("TOGETHER_API_KEY", env_key)
    codeflash_output = default_api_key(arg_key); result = codeflash_output # 232ns -> 305ns (23.9% slower)

# --- Deterministic and Cleanliness Tests ---

def test_deterministic_behavior(monkeypatch):
    # The function should always return the same value for the same input
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key("argument_key") # 270ns -> 274ns (1.46% slower)
    codeflash_output = default_api_key() # 1.55μs -> 1.07μs (44.9% faster)
    monkeypatch.delenv("TOGETHER_API_KEY", raising=False)
    with pytest.raises(error.AuthenticationError):
        default_api_key() # 5.34μs -> 5.23μs (2.20% faster)

def test_no_side_effects(monkeypatch):
    # Ensure function does not modify environment variables
    monkeypatch.setenv("TOGETHER_API_KEY", "env_key")
    codeflash_output = default_api_key(); _ = codeflash_output # 1.55μs -> 1.10μs (40.1% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from together.utils.api_helpers import default_api_key
import pytest

def test_default_api_key():
    with pytest.raises(AuthenticationError, match='TOGETHER_API_KEY\\ not\\ found\\.\\\nPlease\\ set\\ it\\ as\\ an\\ environment\\ variable\\ or\\ set\\ it\\ as\\ together\\.api_key\\\nFind\\ your\\ TOGETHER_API_KEY\\ at\\ https://api\\.together\\.xyz/settings/api\\-keys'):
        default_api_key(api_key='')

def test_default_api_key_2():
    default_api_key(api_key='\x00')