Skip to content

Commit 6fc172e

Browse files
yiyuan-heyiyuan-he
committed
re-use botocore aws region fallback logic as much as possible
1 parent 2c5ffd6 commit 6fc172e

File tree

2 files changed

+114
-28
lines changed
  • aws-opentelemetry-distro

2 files changed

+114
-28
lines changed

‎aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_utils.py‎

Lines changed: 9 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -32,34 +32,29 @@ def is_agent_observability_enabled() -> bool:
3232

3333

3434
def get_aws_region() -> str:
35-
"""Get AWS region from environment or botocore session.
35+
"""Get AWS region using botocore session.
3636
37-
Returns the AWS region in the following priority order:
37+
botocore automatically checks in the following priority order:
3838
1. AWS_REGION environment variable
3939
2. AWS_DEFAULT_REGION environment variable
40-
3. botocore session's region (if botocore is available)
41-
4. None if no region can be determined
42-
"""
43-
# Check AWS environment variables first
44-
region = os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION")
45-
if region:
46-
return region
40+
3. AWS CLI config file (~/.aws/config)
41+
4. EC2 instance metadata service
4742
48-
# Try to get region from botocore if available
49-
# This will automatically check AWS CLI config, instance metadata, etc.
43+
Returns:
44+
The AWS region if found, None otherwise.
45+
"""
5046
if is_installed("botocore"):
5147
try:
5248
from botocore import session # pylint: disable=import-outside-toplevel
5349

5450
botocore_session = session.Session()
55-
if hasattr(botocore_session, "region_name") and botocore_session.region_name:
51+
if botocore_session.region_name:
5652
return botocore_session.region_name
5753
except (ImportError, AttributeError):
5854
# botocore failed to determine region
5955
pass
6056

6157
_logger.warning(
62-
"AWS region not found in environment variables (AWS_REGION, AWS_DEFAULT_REGION) "
63-
"or botocore configuration. Please set AWS_REGION environment variable explicitly."
58+
"AWS region not found. Please set AWS_REGION environment variable or configure " "AWS CLI with 'aws configure'."
6459
)
6560
return None

‎aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_utils.py‎

Lines changed: 105 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# SPDX-License-Identifier: Apache-2.0
33

44
import os
5+
import sys
56
from unittest import TestCase
67
from unittest.mock import MagicMock, patch
78

@@ -19,31 +20,97 @@ def tearDown(self):
1920
os.environ.pop("AWS_REGION", None)
2021
os.environ.pop("AWS_DEFAULT_REGION", None)
2122

22-
def test_get_aws_region_from_aws_region_env(self):
23+
@patch("amazon.opentelemetry.distro._utils.is_installed")
24+
def test_get_aws_region_from_aws_region_env(self, mock_is_installed):
25+
mock_is_installed.return_value = True
26+
27+
# Mock botocore module and session
28+
mock_botocore = MagicMock()
29+
mock_session_instance = MagicMock()
30+
mock_session_instance.region_name = "us-west-2"
31+
mock_botocore.session.Session.return_value = mock_session_instance
32+
sys.modules["botocore"] = mock_botocore
33+
sys.modules["botocore.session"] = mock_botocore.session
34+
2335
os.environ["AWS_REGION"] = "us-west-2"
24-
self.assertEqual(get_aws_region(), "us-west-2")
2536

26-
def test_get_aws_region_from_aws_default_region_env(self):
37+
try:
38+
self.assertEqual(get_aws_region(), "us-west-2")
39+
finally:
40+
# Clean up mock
41+
if "botocore" in sys.modules:
42+
del sys.modules["botocore"]
43+
if "botocore.session" in sys.modules:
44+
del sys.modules["botocore.session"]
45+
46+
@patch("amazon.opentelemetry.distro._utils.is_installed")
47+
def test_get_aws_region_from_aws_default_region_env(self, mock_is_installed):
48+
mock_is_installed.return_value = True
49+
50+
# Mock botocore module and session
51+
mock_botocore = MagicMock()
52+
mock_session_instance = MagicMock()
53+
mock_session_instance.region_name = "eu-central-1"
54+
mock_botocore.session.Session.return_value = mock_session_instance
55+
sys.modules["botocore"] = mock_botocore
56+
sys.modules["botocore.session"] = mock_botocore.session
57+
2758
os.environ["AWS_DEFAULT_REGION"] = "eu-central-1"
28-
self.assertEqual(get_aws_region(), "eu-central-1")
2959

30-
def test_get_aws_region_prefers_aws_region_over_default(self):
60+
try:
61+
self.assertEqual(get_aws_region(), "eu-central-1")
62+
finally:
63+
# Clean up mock
64+
if "botocore" in sys.modules:
65+
del sys.modules["botocore"]
66+
if "botocore.session" in sys.modules:
67+
del sys.modules["botocore.session"]
68+
69+
@patch("amazon.opentelemetry.distro._utils.is_installed")
70+
def test_get_aws_region_prefers_aws_region_over_default(self, mock_is_installed):
71+
mock_is_installed.return_value = True
72+
73+
# Mock botocore module and session
74+
mock_botocore = MagicMock()
75+
mock_session_instance = MagicMock()
76+
mock_session_instance.region_name = "us-east-1"
77+
mock_botocore.session.Session.return_value = mock_session_instance
78+
sys.modules["botocore"] = mock_botocore
79+
sys.modules["botocore.session"] = mock_botocore.session
80+
3181
os.environ["AWS_REGION"] = "us-east-1"
3282
os.environ["AWS_DEFAULT_REGION"] = "eu-west-1"
33-
self.assertEqual(get_aws_region(), "us-east-1")
83+
84+
try:
85+
self.assertEqual(get_aws_region(), "us-east-1")
86+
finally:
87+
# Clean up mock
88+
if "botocore" in sys.modules:
89+
del sys.modules["botocore"]
90+
if "botocore.session" in sys.modules:
91+
del sys.modules["botocore.session"]
3492

3593
@patch("amazon.opentelemetry.distro._utils.is_installed")
36-
@patch("botocore.session.Session")
37-
def test_get_aws_region_from_botocore_session(self, mock_session_class, mock_is_installed):
94+
def test_get_aws_region_from_botocore_session(self, mock_is_installed):
3895
mock_is_installed.return_value = True
3996

40-
mock_session = MagicMock()
41-
mock_session.region_name = "ap-southeast-1"
42-
mock_session_class.return_value = mock_session
97+
# Mock botocore module and session
98+
mock_botocore = MagicMock()
99+
mock_session_instance = MagicMock()
100+
mock_session_instance.region_name = "ap-southeast-1"
101+
mock_botocore.session.Session.return_value = mock_session_instance
102+
sys.modules["botocore"] = mock_botocore
103+
sys.modules["botocore.session"] = mock_botocore.session
43104

44-
result = get_aws_region()
45-
46-
self.assertEqual(result, "ap-southeast-1")
105+
try:
106+
result = get_aws_region()
107+
self.assertEqual(result, "ap-southeast-1")
108+
finally:
109+
# Clean up mock
110+
if "botocore" in sys.modules:
111+
del sys.modules["botocore"]
112+
if "botocore.session" in sys.modules:
113+
del sys.modules["botocore.session"]
47114

48115
@patch("amazon.opentelemetry.distro._utils.is_installed")
49116
@patch("amazon.opentelemetry.distro._utils._logger")
@@ -54,3 +121,27 @@ def test_get_aws_region_returns_none_when_no_region_found(self, mock_logger, moc
54121

55122
self.assertIsNone(result)
56123
mock_logger.warning.assert_called_once()
124+
125+
@patch("amazon.opentelemetry.distro._utils.is_installed")
126+
@patch("amazon.opentelemetry.distro._utils._logger")
127+
def test_get_aws_region_returns_none_when_botocore_has_no_region(self, mock_logger, mock_is_installed):
128+
mock_is_installed.return_value = True
129+
130+
# Mock botocore module with no region
131+
mock_botocore = MagicMock()
132+
mock_session_instance = MagicMock()
133+
mock_session_instance.region_name = None
134+
mock_botocore.session.Session.return_value = mock_session_instance
135+
sys.modules["botocore"] = mock_botocore
136+
sys.modules["botocore.session"] = mock_botocore.session
137+
138+
try:
139+
result = get_aws_region()
140+
self.assertIsNone(result)
141+
mock_logger.warning.assert_called_once()
142+
finally:
143+
# Clean up mock
144+
if "botocore" in sys.modules:
145+
del sys.modules["botocore"]
146+
if "botocore.session" in sys.modules:
147+
del sys.modules["botocore.session"]

0 commit comments

Comments
 (0)