Improve instance metadata region fetcher tests to support IMDSv2

tests/functional/test_clidriver.py

@@ -29,13 +29,13 @@
 
 class RegionCapture(object):
     def __init__(self):
-        self._region = None
+        self.region = None
 
     def __call__(self, request, **kwargs):
         url = request.url
         region = re.match(
             'https://.*?\.(.*?)\.amazonaws\.com', url).groups(1)[0]
-        self._region = region
+        self.region = region
 
 
 class TestSession(BaseCLIDriverTest):
@@ -82,7 +82,7 @@ def test_imds_region_is_used_as_fallback_wo_v2_support(self):
         capture = RegionCapture()
         self.session.register('before-send.ec2.*', capture)
         self.driver.main(['ec2', 'describe-instances'])
-        self.assertEqual(capture._region, 'us-mars-2')
+        self.assertEqual(capture.region, 'us-mars-2')
 
     def test_imds_region_is_used_as_fallback_with_v2_support(self):
         # Remove region override from the environment variables.
@@ -101,7 +101,7 @@ def test_imds_region_is_used_as_fallback_with_v2_support(self):
         capture = RegionCapture()
         self.session.register('before-send.ec2.*', capture)
         self.driver.main(['ec2', 'describe-instances'])
-        self.assertEqual(capture._region, 'us-mars-2')
+        self.assertEqual(capture.region, 'us-mars-2')
 
 
 class TestPlugins(BaseCLIDriverTest):

tests/unit/test_utils.py

@@ -324,6 +324,9 @@ def add_get_region_imds_response(self, region=None):
             region = self._region
         self.add_imds_response(body=region.encode('utf-8'))
 
+    def add_imds_token_response(self):
+        self.add_imds_response(status_code=200, body=b'token')
+
     def add_imds_connection_error(self, exception):
         self._imds_responses.append(exception)
 
@@ -351,6 +354,7 @@ def test_disabling_env_var_not_true(self):
         url = 'https://example.com/'
         env = {'AWS_EC2_METADATA_DISABLED': 'false'}
 
+        self.add_imds_token_response()
         self.add_get_region_imds_response()
 
         fetcher = InstanceMetadataRegionFetcher(base_url=url, env=env)
@@ -361,6 +365,7 @@ def test_disabling_env_var_not_true(self):
 
     def test_includes_user_agent_header(self):
         user_agent = 'my-user-agent'
+        self.add_imds_token_response()
         self.add_get_region_imds_response()
 
         InstanceMetadataRegionFetcher(
@@ -372,6 +377,7 @@ def test_includes_user_agent_header(self):
     def test_non_200_response_for_region_is_retried(self):
         # Response for role name that have a non 200 status code should
         # be retried.
+        self.add_imds_token_response()
         self.add_imds_response(
             status_code=429, body=b'{"message": "Slow down"}')
         self.add_get_region_imds_response()
@@ -393,6 +399,7 @@ def test_empty_response_for_region_is_retried(self):
     def test_non_200_response_is_retried(self):
         # Response for creds that has a 200 status code but has an empty
         # body should be retried.
+        self.add_imds_token_response()
         self.add_imds_response(
             status_code=429, body=b'{"message": "Slow down"}')
         self.add_get_region_imds_response()
@@ -403,6 +410,7 @@ def test_non_200_response_is_retried(self):
 
     def test_http_connection_errors_is_retried(self):
         # Connection related errors should be retried
+        self.add_imds_token_response()
         self.add_imds_connection_error(ConnectionClosedError(endpoint_url=''))
         self.add_get_region_imds_response()
         result = InstanceMetadataRegionFetcher(
@@ -421,6 +429,7 @@ def test_empty_response_is_retried(self):
         self.assertEqual(result, expected_result)
 
     def test_exhaust_retries_on_region_request(self):
+        self.add_imds_token_response()
         self.add_imds_response(status_code=400, body=b'')
         result = InstanceMetadataRegionFetcher(
             num_attempts=1).retrieve_region()