Refactor e2e tests to use assertion helper class

test/e2e/tests/helper.py

@@ -0,0 +1,106 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may
+# not use this file except in compliance with the License. A copy of the
+# License is located at
+#
+#	 http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is distributed
+# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+# express or implied. See the License for the specific language governing
+# permissions and limitations under the License.
+
+"""Helper functions for ec2 tests
+"""
+
+class EC2Validator:
+    def __init__(self, ec2_client):
+        self.ec2_client = ec2_client
+
+    def assert_dhcp_options(self, dhcp_options_id: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_dhcp_options(DhcpOptionsIds=[dhcp_options_id])
+            res_found = len(aws_res["DhcpOptions"]) > 0
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists
+
+    def assert_internet_gateway(self, ig_id: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_internet_gateways(InternetGatewayIds=[ig_id])
+            res_found = len(aws_res["InternetGateways"]) > 0
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists
+
+    def assert_route(self, route_table_id: str, gateway_id: str, origin: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_route_tables(RouteTableIds=[route_table_id])
+            routes = aws_res["RouteTables"][0]["Routes"]
+            for route in routes:
+                if route["Origin"] == origin and route["GatewayId"] == gateway_id:
+                    res_found = True
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists
+
+    def assert_route_table(self, route_table_id: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_route_tables(RouteTableIds=[route_table_id])
+            res_found = len(aws_res["RouteTables"]) > 0
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists
+
+    def assert_security_group(self, sg_id: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_security_groups(GroupIds=[sg_id])
+            res_found = len(aws_res["SecurityGroups"]) > 0
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists
+
+    def assert_subnet(self, subnet_id: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_subnets(SubnetIds=[subnet_id])
+            res_found = len(aws_res["Subnets"]) > 0
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists
+
+    def assert_transit_gateway(self, tgw_id: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_transit_gateways(TransitGatewayIds=[tgw_id])
+            tgw = aws_res["TransitGateways"][0]
+            # TransitGateway may take awhile to be removed server-side, so 
+            # treat 'deleting' and 'deleted' states as resource no longer existing
+            res_found = tgw is not None and tgw['State'] != "deleting" and tgw['State'] != "deleted"
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists
+
+    def assert_vpc(self, vpc_id: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_vpcs(VpcIds=[vpc_id])
+            res_found = len(aws_res["Vpcs"]) > 0
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists
+
+    def assert_vpc_endpoint(self, vpc_endpoint_id: str, exists=True):
+        res_found = False
+        try:
+            aws_res = self.ec2_client.describe_vpc_endpoints(VpcEndpointIds=[vpc_endpoint_id])
+            res_found = len(aws_res["VpcEndpoints"]) > 0
+        except self.ec2_client.exceptions.ClientError:
+            pass
+        assert res_found is exists

test/e2e/tests/test_dhcp_options.py

@@ -22,31 +22,14 @@
 from acktest.k8s import resource as k8s
 from e2e import service_marker, CRD_GROUP, CRD_VERSION, load_ec2_resource
 from e2e.replacement_values import REPLACEMENT_VALUES
+from e2e.tests.helper import EC2Validator
 
 RESOURCE_PLURAL = "dhcpoptions"
 
 DEFAULT_WAIT_AFTER_SECONDS = 5
 CREATE_WAIT_AFTER_SECONDS = 10
 DELETE_WAIT_AFTER_SECONDS = 10
 
-
-def get_dhcp_options(ec2_client, dhcp_options_id: str) -> dict:
-    try:
-        resp = ec2_client.describe_dhcp_options(
-            Filters=[{"Name": "dhcp-options-id", "Values": [dhcp_options_id]}]
-        )
-    except Exception as e:
-        logging.debug(e)
-        return None
-
-    if len(resp["DhcpOptions"]) == 0:
-        return None
-    return resp["DhcpOptions"][0]
-
-
-def dhcp_options_exist(ec2_client, dhcp_options_id: str) -> bool:
-    return get_dhcp_options(ec2_client, dhcp_options_id) is not None
-
 @service_marker
 @pytest.mark.canary
 class TestDhcpOptions:
@@ -84,14 +67,15 @@ def test_create_delete(self, ec2_client):
 
         time.sleep(CREATE_WAIT_AFTER_SECONDS)
 
-        # Check DHCP Options exists
-        assert dhcp_options_exist(ec2_client, resource_id)
+        # Check DHCP Options exists in AWS
+        ec2_validator = EC2Validator(ec2_client)
+        ec2_validator.assert_dhcp_options(resource_id)
 
         # Delete k8s resource
         _, deleted = k8s.delete_custom_resource(ref)
         assert deleted is True
 
         time.sleep(DELETE_WAIT_AFTER_SECONDS)
 
-        # Check DHCP Options doesn't exist
-        assert not dhcp_options_exist(ec2_client, resource_id)
+        # Check DHCP Options no longer exists in AWS
+        ec2_validator.assert_dhcp_options(resource_id, exists=False)

test/e2e/tests/test_internet_gateway.py

@@ -22,30 +22,13 @@
 from acktest.k8s import resource as k8s
 from e2e import service_marker, CRD_GROUP, CRD_VERSION, load_ec2_resource
 from e2e.replacement_values import REPLACEMENT_VALUES
+from e2e.tests.helper import EC2Validator
 
 RESOURCE_PLURAL = "internetgateways"
 
 CREATE_WAIT_AFTER_SECONDS = 10
 DELETE_WAIT_AFTER_SECONDS = 10
 
-
-def get_internet_gateway(ec2_client, ig_id: str) -> dict:
-    try:
-        resp = ec2_client.describe_internet_gateways(
-            Filters=[{"Name": "internet-gateway-id", "Values": [ig_id]}]
-        )
-    except Exception as e:
-        logging.debug(e)
-        return None
-
-    if len(resp["InternetGateways"]) == 0:
-        return None
-    return resp["InternetGateways"][0]
-
-
-def internet_gateway_exists(ec2_client, ig_id: str) -> bool:
-    return get_internet_gateway(ec2_client, ig_id) is not None
-
 @service_marker
 @pytest.mark.canary
 class TestInternetGateway:
@@ -77,16 +60,15 @@ def test_create_delete(self, ec2_client):
 
         time.sleep(CREATE_WAIT_AFTER_SECONDS)
 
-        # Check Internet Gateway exists
-        exists = internet_gateway_exists(ec2_client, resource_id)
-        assert exists
+        # Check Internet Gateway exists in AWS
+        ec2_validator = EC2Validator(ec2_client)
+        ec2_validator.assert_internet_gateway(resource_id)
 
         # Delete k8s resource
         _, deleted = k8s.delete_custom_resource(ref, 2, 5)
         assert deleted is True
 
         time.sleep(DELETE_WAIT_AFTER_SECONDS)
 
-        # Check Internet Gateway doesn't exist
-        exists = internet_gateway_exists(ec2_client, resource_id)
-        assert not exists
+        # Check Internet Gateway no longer exists in AWS
+        ec2_validator.assert_internet_gateway(resource_id, exists=False)

test/e2e/tests/test_route_table.py

@@ -23,51 +23,14 @@
 from e2e import service_marker, CRD_GROUP, CRD_VERSION, load_ec2_resource
 from e2e.replacement_values import REPLACEMENT_VALUES
 from e2e.bootstrap_resources import get_bootstrap_resources
+from e2e.tests.helper import EC2Validator
 
 RESOURCE_PLURAL = "routetables"
 
 DEFAULT_WAIT_AFTER_SECONDS = 5
 CREATE_WAIT_AFTER_SECONDS = 10
 DELETE_WAIT_AFTER_SECONDS = 10
 
-
-def get_route_table(ec2_client, route_table_id: str) -> dict:
-    try:
-        resp = ec2_client.describe_route_tables(
-            Filters=[{"Name": "route-table-id", "Values": [route_table_id]}]
-        )
-    except Exception as e:
-        logging.debug(e)
-        return None
-
-    if len(resp["RouteTables"]) == 0:
-        return None
-    return resp["RouteTables"][0]
-
-
-def route_table_exists(ec2_client, route_table_id: str) -> bool:
-    return get_route_table(ec2_client, route_table_id) is not None
-
-def get_routes(ec2_client, route_table_id: str) -> list:
-    try:
-        resp = ec2_client.describe_route_tables(
-            Filters=[{"Name": "route-table-id", "Values": [route_table_id]}]
-        )
-    except Exception as e:
-        logging.debug(e)
-        return None
-
-    if len(resp["RouteTables"]) == 0:
-        return None
-    return resp["RouteTables"][0]["Routes"]
-
-def route_exists(ec2_client, route_table_id: str, gateway_id: str, origin: str) -> bool:
-    routes = get_routes(ec2_client, route_table_id)
-    for route in routes:
-        if route["Origin"] == origin and route["GatewayId"] == gateway_id:
-            return True
-    return False
-
 @service_marker
 @pytest.mark.canary
 class TestRouteTable:
@@ -107,18 +70,18 @@ def test_create_delete(self, ec2_client):
 
         time.sleep(CREATE_WAIT_AFTER_SECONDS)
 
-        # Check Route Table exists
-        assert route_table_exists(ec2_client, resource_id)
+        # Check Route Table exists in AWS
+        ec2_validator = EC2Validator(ec2_client)
+        ec2_validator.assert_route_table(resource_id)
 
         # Delete k8s resource
         _, deleted = k8s.delete_custom_resource(ref)
         assert deleted is True
 
         time.sleep(DELETE_WAIT_AFTER_SECONDS)
 
-        # Check Route Table doesn't exist
-        exists = route_table_exists(ec2_client, resource_id)
-        assert not exists
+        # Check Route Table no longer exists in AWS
+        ec2_validator.assert_route_table(resource_id, exists=False)
 
 
     def test_terminal_condition(self):
@@ -189,21 +152,16 @@ def test_crud_route(self, ec2_client):
 
         time.sleep(CREATE_WAIT_AFTER_SECONDS)
 
-        # Check Route Table exists
-        assert route_table_exists(ec2_client, resource_id)
-
-        # Check Routes exist (default and desired)
-        routes = get_routes(ec2_client, resource_id)
-        for route in routes:
-            if route["GatewayId"] == "local":
-                default_cidr = route["DestinationCidrBlock"]
-                assert route["Origin"] == "CreateRouteTable"
-            elif route["GatewayId"] == igw_id:
-                assert route["Origin"] == "CreateRoute"
-            else:
-                assert False
+        # Check Route Table exists in AWS
+        ec2_validator = EC2Validator(ec2_client)
+        ec2_validator.assert_route_table(resource_id)
+
+        # Check Routes exist (default and desired) in AWS
+        ec2_validator.assert_route(resource_id, "local", "CreateRouteTable")
+        ec2_validator.assert_route(resource_id, igw_id, "CreateRoute")
 
         # Update Route
+        default_cidr = "10.0.0.0/16"
         updated_cidr = "192.168.1.0/24"
         patch = {"spec": {"routes": [
                     {
@@ -224,15 +182,6 @@ def test_crud_route(self, ec2_client):
         # assert patched state
         resource = k8s.get_resource(ref)
         assert len(resource['status']['routeStatuses']) == 2
-        for route in resource['status']['routeStatuses']:
-            if route["gatewayID"] == "local":
-                assert route_exists(ec2_client, resource_id, "local", "CreateRouteTable")
-            elif route["gatewayID"] == igw_id:
-                # origin and state are set server-side
-                assert route_exists(ec2_client, resource_id, igw_id, "CreateRoute")
-                assert route["state"] == "active"
-            else:
-                assert False
 
         # Delete Route
         patch = {"spec": {"routes": [
@@ -248,12 +197,10 @@ def test_crud_route(self, ec2_client):
 
         resource = k8s.get_resource(ref)
         assert len(resource['spec']['routes']) == 1
-        for route in resource['spec']['routes']:
-            if route["gatewayID"] == "local":
-                assert route_exists(ec2_client, resource_id, "local", "CreateRouteTable")
-            else:
-                assert False
 
+        # Route should no longer exist in AWS (default will remain)
+        ec2_validator.assert_route(resource_id, "local", "CreateRouteTable")
+        ec2_validator.assert_route(resource_id, igw_id, "CreateRoute", exists=False)
 
         # Should not be able to delete default route
         patch = {"spec": {"routes": [
@@ -273,6 +220,5 @@ def test_crud_route(self, ec2_client):
 
         time.sleep(DELETE_WAIT_AFTER_SECONDS)
 
-        # Check Route Table doesn't exist
-        exists = route_table_exists(ec2_client, resource_id)
-        assert not exists
+        # Check Route Table no longer exists in AWS
+        ec2_validator.assert_route_table(resource_id, exists=False)