|
| 1 | +#!/usr/bin/env python |
| 2 | +# - |
| 3 | +# #%L |
| 4 | +# Contrast AI SmartFix |
| 5 | +# %% |
| 6 | +# Copyright (C) 2025 Contrast Security, Inc. |
| 7 | +# %% |
| 8 | + |
| 9 | +# License: Commercial |
| 10 | +# NOTICE: This Software and the patented inventions embodied within may only be |
| 11 | +# used as part of Contrast Security's commercial offerings. Even though it is |
| 12 | +# made available through public repositories, use of this Software is subject to |
| 13 | +# the applicable End User Licensing Agreement found at |
| 14 | +# https://www.contrastsecurity.com/enduser-terms-0317a or as otherwise agreed |
| 15 | +# between Contrast Security and the End User. The Software may not be reverse |
| 16 | +# engineered, modified, repackaged, sold, redistributed or otherwise used in a |
| 17 | +# way not consistent with the End User License Agreement. |
| 18 | +# #L% |
| 19 | +# |
| 20 | + |
| 21 | +""" |
| 22 | +Unit tests for AWS_BEARER_TOKEN_BEDROCK environment variable support. |
| 23 | +
|
| 24 | +This module tests that the AWS_BEARER_TOKEN_BEDROCK environment variable |
| 25 | +is properly handled and available for LiteLLM to use for Bedrock authentication. |
| 26 | +""" |
| 27 | + |
| 28 | +import sys |
| 29 | +import unittest |
| 30 | +import os |
| 31 | + |
| 32 | +# Add project root to path for imports |
| 33 | +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) |
| 34 | + |
| 35 | +# Import config module (need to initialize before other imports) |
| 36 | +from src.config import get_config # noqa: E402 |
| 37 | + |
| 38 | +# Initialize config with testing flag |
| 39 | +_ = get_config(testing=True) |
| 40 | + |
| 41 | + |
| 42 | +class TestAwsBearerTokenBedrock(unittest.TestCase): |
| 43 | + """Test cases for AWS_BEARER_TOKEN_BEDROCK environment variable support.""" |
| 44 | + |
| 45 | + def test_environment_variable_can_be_set(self): |
| 46 | + """Test that AWS_BEARER_TOKEN_BEDROCK environment variable can be set and retrieved.""" |
| 47 | + test_token = "test-bearer-token-12345" |
| 48 | + |
| 49 | + # Set the environment variable |
| 50 | + os.environ['AWS_BEARER_TOKEN_BEDROCK'] = test_token |
| 51 | + |
| 52 | + # Verify it can be retrieved |
| 53 | + self.assertEqual(os.environ.get('AWS_BEARER_TOKEN_BEDROCK'), test_token) |
| 54 | + |
| 55 | + # Clean up |
| 56 | + del os.environ['AWS_BEARER_TOKEN_BEDROCK'] |
| 57 | + |
| 58 | + def test_environment_variable_not_set_returns_none(self): |
| 59 | + """Test that missing AWS_BEARER_TOKEN_BEDROCK returns None.""" |
| 60 | + # Ensure the variable is not set |
| 61 | + if 'AWS_BEARER_TOKEN_BEDROCK' in os.environ: |
| 62 | + del os.environ['AWS_BEARER_TOKEN_BEDROCK'] |
| 63 | + |
| 64 | + # Verify it returns None when not set |
| 65 | + self.assertIsNone(os.environ.get('AWS_BEARER_TOKEN_BEDROCK')) |
| 66 | + |
| 67 | + def test_bearer_token_precedence_over_iam(self): |
| 68 | + """Test that AWS_BEARER_TOKEN_BEDROCK can coexist with IAM credentials.""" |
| 69 | + # Set both bearer token and IAM credentials |
| 70 | + os.environ['AWS_BEARER_TOKEN_BEDROCK'] = "test-bearer-token" |
| 71 | + os.environ['AWS_ACCESS_KEY_ID'] = "test-access-key" |
| 72 | + os.environ['AWS_SECRET_ACCESS_KEY'] = "test-secret-key" |
| 73 | + |
| 74 | + try: |
| 75 | + # Verify both are set (LiteLLM will determine which to use based on its own logic) |
| 76 | + self.assertIsNotNone(os.environ.get('AWS_BEARER_TOKEN_BEDROCK')) |
| 77 | + self.assertIsNotNone(os.environ.get('AWS_ACCESS_KEY_ID')) |
| 78 | + self.assertIsNotNone(os.environ.get('AWS_SECRET_ACCESS_KEY')) |
| 79 | + |
| 80 | + finally: |
| 81 | + # Clean up |
| 82 | + for key in ['AWS_BEARER_TOKEN_BEDROCK', 'AWS_ACCESS_KEY_ID', 'AWS_SECRET_ACCESS_KEY']: |
| 83 | + if key in os.environ: |
| 84 | + del os.environ[key] |
| 85 | + |
| 86 | + def test_empty_bearer_token_is_ignored(self): |
| 87 | + """Test that an empty AWS_BEARER_TOKEN_BEDROCK value is handled gracefully.""" |
| 88 | + # Set an empty bearer token |
| 89 | + os.environ['AWS_BEARER_TOKEN_BEDROCK'] = "" |
| 90 | + |
| 91 | + try: |
| 92 | + # Verify it's set but empty |
| 93 | + self.assertEqual(os.environ.get('AWS_BEARER_TOKEN_BEDROCK'), "") |
| 94 | + |
| 95 | + # An empty token should be falsy |
| 96 | + self.assertFalse(os.environ.get('AWS_BEARER_TOKEN_BEDROCK')) |
| 97 | + |
| 98 | + finally: |
| 99 | + # Clean up |
| 100 | + if 'AWS_BEARER_TOKEN_BEDROCK' in os.environ: |
| 101 | + del os.environ['AWS_BEARER_TOKEN_BEDROCK'] |
| 102 | + |
| 103 | + def test_aws_region_name_can_be_set(self): |
| 104 | + """Test that AWS_REGION_NAME environment variable can be set and retrieved.""" |
| 105 | + test_region = "us-east-1" |
| 106 | + |
| 107 | + # Set the environment variable |
| 108 | + os.environ['AWS_REGION_NAME'] = test_region |
| 109 | + |
| 110 | + try: |
| 111 | + # Verify it can be retrieved |
| 112 | + self.assertEqual(os.environ.get('AWS_REGION_NAME'), test_region) |
| 113 | + |
| 114 | + finally: |
| 115 | + # Clean up |
| 116 | + if 'AWS_REGION_NAME' in os.environ: |
| 117 | + del os.environ['AWS_REGION_NAME'] |
| 118 | + |
| 119 | + def test_bearer_token_and_region_together(self): |
| 120 | + """Test that AWS_BEARER_TOKEN_BEDROCK and AWS_REGION_NAME can be used together.""" |
| 121 | + test_token = "test-bearer-token-abc123" |
| 122 | + test_region = "us-west-2" |
| 123 | + |
| 124 | + # Set both environment variables |
| 125 | + os.environ['AWS_BEARER_TOKEN_BEDROCK'] = test_token |
| 126 | + os.environ['AWS_REGION_NAME'] = test_region |
| 127 | + |
| 128 | + try: |
| 129 | + # Verify both are set correctly |
| 130 | + self.assertEqual(os.environ.get('AWS_BEARER_TOKEN_BEDROCK'), test_token) |
| 131 | + self.assertEqual(os.environ.get('AWS_REGION_NAME'), test_region) |
| 132 | + |
| 133 | + finally: |
| 134 | + # Clean up |
| 135 | + for key in ['AWS_BEARER_TOKEN_BEDROCK', 'AWS_REGION_NAME']: |
| 136 | + if key in os.environ: |
| 137 | + del os.environ[key] |
| 138 | + |
| 139 | + |
| 140 | +if __name__ == '__main__': |
| 141 | + unittest.main() |
0 commit comments