Add tests for AwsInstance

and change urllib.request.urlopen() to requests.get() for easier testing; functionality doesn't change
2021-02-06 19:19:08 +05:30 · 2021-02-06 19:19:08 +05:30 · 016d886781
parent eed5ea1337
commit 016d886781
2 changed files with 266 additions and 10 deletions
--- a/monkey/common/cloud/aws/aws_instance.py
+++ b/monkey/common/cloud/aws/aws_instance.py
@ -1,8 +1,7 @@
 import json
 import logging
 import re
-import urllib.error
-import urllib.request
+import requests

 from common.cloud.environment_names import Environment
 from common.cloud.instance import CloudInstance
@ -33,19 +32,16 @@ class AwsInstance(CloudInstance):
        self.account_id = None

        try:
-            self.instance_id = urllib.request.urlopen(
-                AWS_LATEST_METADATA_URI_PREFIX + 'meta-data/instance-id', timeout=2).read().decode()
+            self.instance_id = requests.get(AWS_LATEST_METADATA_URI_PREFIX + 'meta-data/instance-id', timeout=2).text
            self.region = self._parse_region(
-                urllib.request.urlopen(
-                    AWS_LATEST_METADATA_URI_PREFIX + 'meta-data/placement/availability-zone').read().decode())
-        except (urllib.error.URLError, IOError) as e:
+                requests.get(AWS_LATEST_METADATA_URI_PREFIX + 'meta-data/placement/availability-zone').text)
+        except (requests.RequestException, IOError) as e:
            logger.debug("Failed init of AwsInstance while getting metadata: {}".format(e))

        try:
            self.account_id = self._extract_account_id(
-                urllib.request.urlopen(
-                    AWS_LATEST_METADATA_URI_PREFIX + 'dynamic/instance-identity/document', timeout=2).read().decode())
-        except (urllib.error.URLError, json.decoder.JSONDecodeError, IOError) as e:
+                requests.get(AWS_LATEST_METADATA_URI_PREFIX + 'dynamic/instance-identity/document', timeout=2).text)
+        except (requests.RequestException, json.decoder.JSONDecodeError, IOError) as e:
            logger.debug("Failed init of AwsInstance while getting dynamic instance data: {}".format(e))

    @staticmethod
--- a/monkey/common/cloud/aws/test_aws_instance.py
+++ b/monkey/common/cloud/aws/test_aws_instance.py
@ -0,0 +1,260 @@
+import json
+
+import pytest
+import requests
+import requests_mock
+
+from common.cloud.aws.aws_instance import (AWS_LATEST_METADATA_URI_PREFIX,
+                                           AwsInstance)
+from common.cloud.environment_names import Environment
+
+
+INSTANCE_ID_RESPONSE = 'i-1234567890abcdef0'
+
+AVAILABILITY_ZONE_RESPONSE = 'us-west-2b'
+
+# from https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/instance-identity-documents.html
+INSTANCE_IDENTITY_DOCUMENT_RESPONSE = """
+{
+    "devpayProductCodes": null,
+    "marketplaceProductCodes": ["1abc2defghijklm3nopqrs4tu"],
+    "availabilityZone": "us-west-2b",
+    "privateIp": "10.158.112.84",
+    "version": "2017-09-30",
+    "instanceId": "i-1234567890abcdef0",
+    "billingProducts": null,
+    "instanceType": "t2.micro",
+    "accountId": "123456789012",
+    "imageId": "ami-5fb8c835",
+    "pendingTime": "2016-11-19T16:32:11Z",
+    "architecture": "x86_64",
+    "kernelId": null,
+    "ramdiskId": null,
+    "region": "us-west-2"
+}
+"""
+
+
+EXPECTED_INSTANCE_ID = 'i-1234567890abcdef0'
+
+EXPECTED_REGION = 'us-west-2'
+
+EXPECTED_ACCOUNT_ID = '123456789012'
+
+
+def get_test_aws_instance(text={'instance_id': None,
+                                'region': None,
+                                'account_id': None},
+                          exception={'instance_id': None,
+                                     'region': None,
+                                     'account_id': None}):
+    with requests_mock.Mocker() as m:
+        # request made to get instance_id
+        url = f'{AWS_LATEST_METADATA_URI_PREFIX}meta-data/instance-id'
+        m.get(url, text=text['instance_id']) if text['instance_id'] else m.get(
+            url, exc=exception['instance_id'])
+
+        # request made to get region
+        url = f'{AWS_LATEST_METADATA_URI_PREFIX}meta-data/placement/availability-zone'
+        m.get(url, text=text['region']) if text['region'] else m.get(
+            url, exc=exception['region'])
+
+        # request made to get account_id
+        url = f'{AWS_LATEST_METADATA_URI_PREFIX}dynamic/instance-identity/document'
+        m.get(url, text=text['account_id']) if text['account_id'] else m.get(
+            url, exc=exception['account_id'])
+
+        test_aws_instance_object = AwsInstance()
+        return test_aws_instance_object
+
+
+# all good data
+@pytest.fixture
+def good_data_mock_instance():
+    return get_test_aws_instance(text={'instance_id': INSTANCE_ID_RESPONSE,
+                                       'region': AVAILABILITY_ZONE_RESPONSE,
+                                       'account_id': INSTANCE_IDENTITY_DOCUMENT_RESPONSE})
+
+
+def test_is_instance_good_data(good_data_mock_instance):
+    assert good_data_mock_instance.is_instance()
+
+
+def test_get_cloud_provider_name_good_data(good_data_mock_instance):
+    assert good_data_mock_instance.get_cloud_provider_name() == Environment.AWS
+
+
+def test_get_instance_id_good_data(good_data_mock_instance):
+    assert good_data_mock_instance.get_instance_id() == EXPECTED_INSTANCE_ID
+
+
+def test_get_region_good_data(good_data_mock_instance):
+    assert good_data_mock_instance.get_region() == EXPECTED_REGION
+
+
+def test_get_account_id_good_data(good_data_mock_instance):
+    assert good_data_mock_instance.get_account_id() == EXPECTED_ACCOUNT_ID
+
+
+# 'region' bad data
+@pytest.fixture
+def bad_data_mock_instance_1():
+    return get_test_aws_instance(text={'instance_id': INSTANCE_ID_RESPONSE,
+                                       'region': 'in-a-different-world',
+                                       'account_id': INSTANCE_IDENTITY_DOCUMENT_RESPONSE})
+
+
+def test_is_instance_bad_data_1(bad_data_mock_instance_1):
+    assert bad_data_mock_instance_1.is_instance()
+
+
+def test_get_cloud_provider_name_bad_data_1(bad_data_mock_instance_1):
+    assert bad_data_mock_instance_1.get_cloud_provider_name() == Environment.AWS
+
+
+def test_get_instance_id_bad_data_1(bad_data_mock_instance_1):
+    assert bad_data_mock_instance_1.get_instance_id() == EXPECTED_INSTANCE_ID
+
+
+def test_get_region_bad_data_1(bad_data_mock_instance_1):
+    assert bad_data_mock_instance_1.get_region() is None
+
+
+def test_get_account_id_bad_data_1(bad_data_mock_instance_1):
+    assert bad_data_mock_instance_1.get_account_id() == EXPECTED_ACCOUNT_ID
+
+
+# 'account_id' bad data
+@pytest.fixture
+def bad_data_mock_instance_2():
+    return get_test_aws_instance(text={'instance_id': INSTANCE_ID_RESPONSE,
+                                       'region': AVAILABILITY_ZONE_RESPONSE,
+                                       'account_id': 'who-am-i'})
+
+
+def test_is_instance_bad_data_2(bad_data_mock_instance_2):
+    assert bad_data_mock_instance_2.is_instance()
+
+
+def test_get_cloud_provider_name_bad_data_2(bad_data_mock_instance_2):
+    assert bad_data_mock_instance_2.get_cloud_provider_name() == Environment.AWS
+
+
+def test_get_instance_id_bad_data_2(bad_data_mock_instance_2):
+    assert bad_data_mock_instance_2.get_instance_id() == EXPECTED_INSTANCE_ID
+
+
+def test_get_region_bad_data_2(bad_data_mock_instance_2):
+    assert bad_data_mock_instance_2.get_region() == EXPECTED_REGION
+
+
+def test_get_account_id_bad_data_2(bad_data_mock_instance_2):
+    assert bad_data_mock_instance_2.get_account_id() is None
+
+
+# 'instance_id' bad requests
+@pytest.fixture
+def bad_request_mock_instance_1(instance_id_exception):
+    return get_test_aws_instance(text={'instance_id': None,
+                                       'region': AVAILABILITY_ZONE_RESPONSE,
+                                       'account_id': INSTANCE_IDENTITY_DOCUMENT_RESPONSE},
+                                 exception={'instance_id': instance_id_exception,
+                                            'region': None,
+                                            'account_id': None})
+
+
+@pytest.mark.parametrize('instance_id_exception', [requests.RequestException, IOError])
+def test_is_instance_bad_request_1(bad_request_mock_instance_1):
+    assert bad_request_mock_instance_1.is_instance() is False
+
+
+@pytest.mark.parametrize('instance_id_exception', [requests.RequestException, IOError])
+def test_get_cloud_provider_name_bad_request_1(bad_request_mock_instance_1):
+    assert bad_request_mock_instance_1.get_cloud_provider_name() == Environment.AWS
+
+
+@pytest.mark.parametrize('instance_id_exception', [requests.RequestException, IOError])
+def test_get_instance_id_bad_request_1(bad_request_mock_instance_1):
+    assert bad_request_mock_instance_1.get_instance_id() is None
+
+
+@pytest.mark.parametrize('instance_id_exception', [requests.RequestException, IOError])
+def test_get_region_bad_request_1(bad_request_mock_instance_1):
+    assert bad_request_mock_instance_1.get_region() is None
+
+
+@pytest.mark.parametrize('instance_id_exception', [requests.RequestException, IOError])
+def test_get_account_id_bad_request_1(bad_request_mock_instance_1):
+    assert bad_request_mock_instance_1.get_account_id() == EXPECTED_ACCOUNT_ID
+
+
+# 'region' bad requests
+@pytest.fixture
+def bad_request_mock_instance_2(region_exception):
+    return get_test_aws_instance(text={'instance_id': INSTANCE_ID_RESPONSE,
+                                       'region': None,
+                                       'account_id': INSTANCE_IDENTITY_DOCUMENT_RESPONSE},
+                                 exception={'instance_id': None,
+                                            'region': region_exception,
+                                            'account_id': None})
+
+
+@pytest.mark.parametrize('region_exception', [requests.RequestException, IOError])
+def test_is_instance_bad_request_2(bad_request_mock_instance_2):
+    assert bad_request_mock_instance_2.is_instance()
+
+
+@pytest.mark.parametrize('region_exception', [requests.RequestException, IOError])
+def test_get_cloud_provider_name_bad_request_2(bad_request_mock_instance_2):
+    assert bad_request_mock_instance_2.get_cloud_provider_name() == Environment.AWS
+
+
+@pytest.mark.parametrize('region_exception', [requests.RequestException, IOError])
+def test_get_instance_id_bad_request_2(bad_request_mock_instance_2):
+    assert bad_request_mock_instance_2.get_instance_id() == EXPECTED_INSTANCE_ID
+
+
+@pytest.mark.parametrize('region_exception', [requests.RequestException, IOError])
+def test_get_region_bad_request_2(bad_request_mock_instance_2):
+    assert bad_request_mock_instance_2.get_region() is None
+
+
+@pytest.mark.parametrize('region_exception', [requests.RequestException, IOError])
+def test_get_account_id_bad_request_2(bad_request_mock_instance_2):
+    assert bad_request_mock_instance_2.get_account_id() == EXPECTED_ACCOUNT_ID
+
+
+# 'account_id' bad requests
+@pytest.fixture
+def bad_request_mock_instance_3(account_id_exception):
+    return get_test_aws_instance(text={'instance_id': INSTANCE_ID_RESPONSE,
+                                       'region': AVAILABILITY_ZONE_RESPONSE,
+                                       'account_id': None},
+                                 exception={'instance_id': None,
+                                            'region': None,
+                                            'account_id': account_id_exception})
+
+
+@pytest.mark.parametrize('account_id_exception', [requests.RequestException, IOError])
+def test_is_instance_bad_request_3(bad_request_mock_instance_3):
+    assert bad_request_mock_instance_3.is_instance()
+
+
+@pytest.mark.parametrize('account_id_exception', [requests.RequestException, IOError])
+def test_get_cloud_provider_name_bad_request_3(bad_request_mock_instance_3):
+    assert bad_request_mock_instance_3.get_cloud_provider_name() == Environment.AWS
+
+
+@pytest.mark.parametrize('account_id_exception', [requests.RequestException, IOError])
+def test_get_instance_id_bad_request_3(bad_request_mock_instance_3):
+    assert bad_request_mock_instance_3.get_instance_id() == EXPECTED_INSTANCE_ID
+
+
+@pytest.mark.parametrize('account_id_exception', [requests.RequestException, IOError])
+def test_get_region_bad_request_3(bad_request_mock_instance_3):
+    assert bad_request_mock_instance_3.get_region() == EXPECTED_REGION
+
+
+@pytest.mark.parametrize('account_id_exception', [requests.RequestException, IOError])
+def test_get_account_id_bad_request_3(bad_request_mock_instance_3):
+    assert bad_request_mock_instance_3.get_account_id() is None