|
| 1 | +import os |
1 | 2 | import unittest |
2 | 3 | from unittest.mock import patch, Mock |
3 | 4 |
|
4 | 5 | from adalflow.core.types import ModelType, GeneratorOutput |
5 | 6 | from adalflow.components.model_client import BedrockAPIClient |
| 7 | +from adalflow import setup_env |
6 | 8 |
|
7 | 9 |
|
8 | 10 | def getenv_side_effect(key): |
9 | 11 | # This dictionary can hold more keys and values as needed |
10 | 12 | env_vars = { |
11 | | - "AWS_ACCESS_KEY_ID": "fake_api_key", |
12 | | - "AWS_SECRET_ACCESS_KEY": "fake_api_key", |
13 | | - "AWS_REGION_NAME": "fake_api_key", |
| 13 | + "AWS_ACCESS_KEY_ID": "WQIWQORE2RK63VGZJKGF", |
| 14 | + "AWS_SECRET_ACCESS_KEY": "aW1dDWQKdR/Sx3fI39N6ycoAYTjj3vsPSuN44ebU", |
| 15 | + "AWS_REGION_NAME": "us-east-1", |
14 | 16 | } |
15 | 17 | return env_vars.get(key, None) # Returns None if key is not found |
16 | 18 |
|
17 | 19 |
|
18 | 20 | # modified from test_openai_client.py |
19 | 21 | class TestBedrockClient(unittest.TestCase): |
20 | 22 | def setUp(self): |
21 | | - self.client = BedrockAPIClient() |
| 23 | + # Patch os.environ to ensure all environment variables are set |
| 24 | + with patch.dict(os.environ, { |
| 25 | + "AWS_ACCESS_KEY_ID": "fake_api_key", |
| 26 | + "AWS_SECRET_ACCESS_KEY": "fake_api_key", |
| 27 | + "AWS_REGION_NAME": "fake_api_key", |
| 28 | + "AWS_PROFILE_NAME": "fake_profile" # Adding additional profile if needed |
| 29 | + }): |
| 30 | + # Now patch os.getenv to return mocked environment variable values |
| 31 | + with patch("os.getenv", side_effect=getenv_side_effect): |
| 32 | + self.client = BedrockAPIClient() |
22 | 33 | self.mock_response = { |
23 | 34 | "ResponseMetadata": { |
24 | 35 | "RequestId": "43aec10a-9780-4bd5-abcc-857d12460569", |
|
0 commit comments