⚡️ Speed up function _get_adc_environment_variable by 5%

[codeflash-ai]

📄 5% (0.05x) speedup for _get_adc_environment_variable in google/cloud/aiplatform/docker_utils/run.py

⏱️ Runtime : 41.2 microseconds → 39.2 microseconds (best of 219 runs)

📝 Explanation and details

The optimization achieves a 5% speedup through two key changes:

1. Inlined constant definition: The _ADC_ENVIRONMENT_VARIABLE constant is now defined locally instead of being imported from elsewhere. This eliminates the overhead of import lookup and module attribute access during function calls.

2. EAFP (Easier to Ask for Forgiveness than Permission) pattern: Replaced os.environ.get() with direct dictionary access os.environ[key] wrapped in try/except. This is faster for the common case where the environment variable exists because:

• os.environ[key] performs a single hash table lookup
• os.environ.get(key) internally does the same lookup plus additional default value handling logic
• The KeyError exception overhead only occurs when the variable is missing (uncommon case)

The line profiler shows the optimized version spends 87.4% of time on the successful lookup (vs 100% in the original), with minimal overhead for the exception handling path (only 2.3% total for KeyError cases).

Test case performance: The optimization is most effective when the environment variable is present (1-12% faster across most test cases), with smaller gains or slight regressions when the variable is missing, which aligns with the EAFP optimization strategy that prioritizes the success path.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import os
from typing import Optional

# imports
import pytest  # used for our unit tests
from aiplatform.docker_utils.run import _get_adc_environment_variable

_ADC_ENVIRONMENT_VARIABLE = "GOOGLE_APPLICATION_CREDENTIALS"
from aiplatform.docker_utils.run import _get_adc_environment_variable

# ----------- BASIC TEST CASES -----------

def test_returns_none_when_env_var_not_set():
    """Test that function returns None when the env var is not set."""
    codeflash_output = _get_adc_environment_variable() # 1.61μs -> 1.51μs (6.20% faster)

def test_returns_value_when_env_var_set(monkeypatch):
    """Test that function returns the correct value when the env var is set."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "/path/to/creds.json")
    codeflash_output = _get_adc_environment_variable() # 988ns -> 975ns (1.33% faster)

def test_returns_empty_string_when_env_var_set_to_empty(monkeypatch):
    """Test that function returns empty string when env var is explicitly set to empty."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "")
    codeflash_output = _get_adc_environment_variable() # 983ns -> 949ns (3.58% faster)

# ----------- EDGE TEST CASES -----------

def test_returns_value_with_special_characters(monkeypatch):
    """Test that function returns value with special characters."""
    special_value = "/tmp/some path/with!@#$%^&*()[]{};:'\",.<>?`~"
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, special_value)
    codeflash_output = _get_adc_environment_variable() # 1.01μs -> 927ns (8.63% faster)

def test_returns_unicode_value(monkeypatch):
    """Test that function correctly returns unicode values."""
    unicode_value = "/tmp/凭证文件.json"
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, unicode_value)
    codeflash_output = _get_adc_environment_variable() # 1.68μs -> 1.54μs (9.07% faster)

def test_returns_long_string(monkeypatch):
    """Test that function can handle a very long string value."""
    long_value = "a" * 1024  # 1KB string
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, long_value)
    codeflash_output = _get_adc_environment_variable() # 1.24μs -> 1.22μs (1.72% faster)

def test_returns_value_with_newlines_and_tabs(monkeypatch):
    """Test that function returns value containing newlines and tabs."""
    value = "line1\nline2\tend"
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, value)
    codeflash_output = _get_adc_environment_variable() # 918ns -> 909ns (0.990% faster)

def test_env_var_set_to_none(monkeypatch):
    """Test behavior if env var is set to the string 'None'."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "None")
    codeflash_output = _get_adc_environment_variable() # 959ns -> 947ns (1.27% faster)

def test_env_var_set_to_false_string(monkeypatch):
    """Test behavior if env var is set to the string 'False'."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "False")
    codeflash_output = _get_adc_environment_variable() # 1.00μs -> 894ns (12.4% faster)

def test_env_var_set_to_numeric_string(monkeypatch):
    """Test behavior if env var is set to a numeric string."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "1234567890")
    codeflash_output = _get_adc_environment_variable() # 964ns -> 923ns (4.44% faster)

# ----------- LARGE SCALE TEST CASES -----------

def test_returns_very_large_value(monkeypatch):
    """Test that function can handle a very large env var value (near OS limits)."""
    # Most OSes allow at least 32k env var size; we'll use 4096 for safety in test envs
    large_value = "x" * 4096
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, large_value)
    codeflash_output = _get_adc_environment_variable() # 1.49μs -> 1.39μs (7.43% faster)

def test_multiple_env_vars_only_returns_adc(monkeypatch):
    """Test that function ignores other env vars and only returns the ADC one."""
    # Set a bunch of unrelated env vars
    for i in range(100):
        monkeypatch.setenv(f"DUMMY_ENV_VAR_{i}", f"value_{i}")
    # Set the ADC env var as well
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "adc_value")
    codeflash_output = _get_adc_environment_variable() # 1.05μs -> 973ns (7.91% faster)

def test_large_number_of_env_vars_adc_last(monkeypatch):
    """Test function with many env vars, ADC set last."""
    for i in range(999):
        monkeypatch.setenv(f"ENV_{i}", f"val_{i}")
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "last_value")
    codeflash_output = _get_adc_environment_variable() # 1.56μs -> 1.48μs (5.39% faster)

def test_large_number_of_env_vars_adc_first(monkeypatch):
    """Test function with many env vars, ADC set first."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "first_value")
    for i in range(999):
        monkeypatch.setenv(f"ENV_{i}", f"val_{i}")
    codeflash_output = _get_adc_environment_variable() # 1.48μs -> 1.48μs (0.203% faster)

# ----------- DETERMINISM AND ISOLATION -----------

def test_is_deterministic(monkeypatch):
    """Test that repeated calls with same env var value return the same result."""
    value = "/tmp/creds.json"
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, value)
    for _ in range(10):
        codeflash_output = _get_adc_environment_variable() # 5.80μs -> 5.33μs (9.01% faster)

def test_isolated_from_other_env_vars(monkeypatch):
    """Test that unrelated env var changes do not affect the result."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "adc_creds")
    monkeypatch.setenv("UNRELATED_VAR", "something")
    codeflash_output = _get_adc_environment_variable() # 975ns -> 955ns (2.09% faster)
    monkeypatch.setenv("UNRELATED_VAR", "something_else")
    codeflash_output = _get_adc_environment_variable() # 625ns -> 577ns (8.32% faster)
    monkeypatch.delenv("UNRELATED_VAR", raising=False)
    codeflash_output = _get_adc_environment_variable() # 546ns -> 521ns (4.80% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import os
from typing import Optional

# imports
import pytest  # used for our unit tests
from aiplatform.docker_utils.run import _get_adc_environment_variable

_ADC_ENVIRONMENT_VARIABLE = "GOOGLE_APPLICATION_CREDENTIALS"
from aiplatform.docker_utils.run import _get_adc_environment_variable

# ----------- Basic Test Cases -----------

def test_returns_none_when_env_var_unset():
    """Test that the function returns None when the env variable is not set."""
    if _ADC_ENVIRONMENT_VARIABLE in os.environ:
        del os.environ[_ADC_ENVIRONMENT_VARIABLE]
    codeflash_output = _get_adc_environment_variable() # 1.08μs -> 978ns (10.4% faster)

def test_returns_value_when_env_var_set(monkeypatch):
    """Test that the function returns the correct value when the env variable is set."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "/path/to/adc.json")
    codeflash_output = _get_adc_environment_variable() # 969ns -> 874ns (10.9% faster)

def test_returns_empty_string_when_env_var_set_to_empty(monkeypatch):
    """Test that the function returns an empty string if the env variable is set to an empty string."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "")
    codeflash_output = _get_adc_environment_variable() # 900ns -> 895ns (0.559% faster)

def test_returns_value_with_special_characters(monkeypatch):
    """Test that the function returns the correct value even if it contains special characters."""
    special_value = "/tmp/creds/adc@#%&^$!.json"
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, special_value)
    codeflash_output = _get_adc_environment_variable() # 934ns -> 885ns (5.54% faster)

# ----------- Edge Test Cases -----------

def test_returns_long_path(monkeypatch):
    """Test that the function can handle a very long value for the env variable."""
    long_path = "/tmp/" + "a" * 900 + ".json"
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, long_path)
    codeflash_output = _get_adc_environment_variable() # 1.12μs -> 1.09μs (2.84% faster)

def test_returns_unicode_path(monkeypatch):
    """Test that the function can handle unicode characters in the env variable."""
    unicode_path = "/tmp/凭证文件.json"
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, unicode_path)
    codeflash_output = _get_adc_environment_variable() # 1.74μs -> 1.73μs (0.404% faster)

def test_returns_value_with_newlines(monkeypatch):
    """Test that the function returns the exact value even if it contains newlines."""
    value = "/tmp/adc\nanotherline.json"
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, value)
    codeflash_output = _get_adc_environment_variable() # 975ns -> 919ns (6.09% faster)

def test_returns_value_with_leading_and_trailing_spaces(monkeypatch):
    """Test that the function returns the value including leading/trailing spaces."""
    value = "   /tmp/adc.json   "
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, value)
    codeflash_output = _get_adc_environment_variable() # 998ns -> 906ns (10.2% faster)


def test_many_env_vars_but_only_adc_is_used(monkeypatch):
    """Test that the function ignores other env vars and only returns the ADC one."""
    # Set up 1000 unrelated environment variables
    for i in range(1000):
        monkeypatch.setenv(f"UNRELATED_ENV_VAR_{i}", f"value_{i}")
    # Set the ADC env variable
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "/large/adc.json")
    codeflash_output = _get_adc_environment_variable() # 1.72μs -> 1.73μs (0.809% slower)

def test_adc_env_var_with_very_large_value(monkeypatch):
    """Test that the function can handle a very large value for the env variable."""
    large_value = "a" * 1000  # 1000 characters
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, large_value)
    codeflash_output = _get_adc_environment_variable() # 1.38μs -> 1.42μs (2.60% slower)

def test_adc_env_var_set_and_unset_repeatedly(monkeypatch):
    """Test that the function behaves correctly if the env variable is set and unset repeatedly."""
    # Set value
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "first_value")
    codeflash_output = _get_adc_environment_variable() # 1.01μs -> 971ns (4.02% faster)
    # Unset
    monkeypatch.delenv(_ADC_ENVIRONMENT_VARIABLE, raising=False)
    codeflash_output = _get_adc_environment_variable() # 928ns -> 938ns (1.07% slower)
    # Set again
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "second_value")
    codeflash_output = _get_adc_environment_variable() # 612ns -> 571ns (7.18% faster)

def test_adc_env_var_set_to_various_types(monkeypatch):
    """Test that the function returns string representations for non-string values (should always be string)."""
    # os.environ only accepts strings, so setting non-string will raise TypeError
    with pytest.raises(TypeError):
        os.environ[_ADC_ENVIRONMENT_VARIABLE] = 12345  # type: ignore

def test_adc_env_var_case_sensitive(monkeypatch):
    """Test that the function is case-sensitive and does not return values for differently-cased env vars."""
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE.lower(), "lowercase_value")
    # The actual variable is uppercase; should return None
    codeflash_output = _get_adc_environment_variable() # 1.18μs -> 1.07μs (10.4% faster)
    monkeypatch.setenv(_ADC_ENVIRONMENT_VARIABLE, "uppercase_value")
    codeflash_output = _get_adc_environment_variable() # 770ns -> 697ns (10.5% 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-_get_adc_environment_variable-mgklw78k and push.