refactor: Move logic for getting regions from the config to common.py

Currently, the `region` key can be configured in the image
configuration. However, sibce `sns` also support the `region`
key to send a notifications to specific regions, this logic
can be resued to get the sns regions.

Move the logic to common.py and rename to `_get_regions` as
a more general function for retrieving regions from
different keys.

Author: JessicaJang

awspub/common.py

@@ -1,4 +1,10 @@
-from typing import Tuple
+import logging
+from typing import List, Tuple
+
+import boto3
+from mypy_boto3_ec2.client import EC2Client
+
+logger = logging.getLogger(__name__)
 
 
 def _split_partition(val: str) -> Tuple[str, str]:
@@ -16,3 +22,43 @@ def _split_partition(val: str) -> Tuple[str, str]:
         partition = "aws"
         resource = val
     return partition, resource
+
+
+def _get_regions(region_to_query: str, regions_allowlist: List[str]) -> List[str]:
+    """
+    Get a list of region names querying the `region_to_query` for all regions and
+    then filtering by `regions_allowlist`.
+    If no `regions_allowlist` is given, all queried regions are returned for the
+    current partition.
+    If `regions_allowlist` is given, all regions from that list are returned if
+    the listed region exist in the current partition.
+    Eg. `us-east-1` listed in `regions_allowlist` won't be returned if the current
+    partition is `aws-cn`.
+    :param region_to_query: region name of current partition
+    :type region_to_query: str
+    :praram regions_allowlist: list of regions in config file
+    :type regions_allowlist: List[str]
+    :return: list of regions names
+    :rtype: List[str]
+    """
+
+    # get all available regions
+    ec2client: EC2Client = boto3.client("ec2", region_name=region_to_query)
+    resp = ec2client.describe_regions()
+    ec2_regions_all = [r["RegionName"] for r in resp["Regions"]]
+
+    if regions_allowlist:
+        # filter out regions that are not available in the current partition
+        regions_allowlist_set = set(regions_allowlist)
+        ec2_regions_all_set = set(ec2_regions_all)
+        regions = list(regions_allowlist_set.intersection(ec2_regions_all_set))
+        diff = regions_allowlist_set.difference(ec2_regions_all_set)
+        if diff:
+            logger.warning(
+                f"regions {diff} listed in regions allowlist are not available in the current partition."
+                " Ignoring those."
+            )
+    else:
+        regions = ec2_regions_all
+
+    return regions

awspub/image.py

@@ -9,7 +9,7 @@
 from mypy_boto3_ssm import SSMClient
 
 from awspub import exceptions
-from awspub.common import _split_partition
+from awspub.common import _get_regions, _split_partition
 from awspub.context import Context
 from awspub.image_marketplace import ImageMarketplace
 from awspub.s3 import S3
@@ -121,28 +121,11 @@ def snapshot_name(self) -> str:
     @property
     def image_regions(self) -> List[str]:
         """
-        Get the image regions. Either configured in the image configuration
-        or all available regions.
-        If a region is listed that is not available in the currently used partition,
-        that region will be ignored (eg. having us-east-1 configured but running in the aws-cn
-        partition doesn't include us-east-1 here).
+        Get the image regions.
         """
         if not self._image_regions_cached:
-            # get all available regions
-            ec2client: EC2Client = boto3.client("ec2", region_name=self._s3.bucket_region)
-            resp = ec2client.describe_regions()
-            image_regions_all = [r["RegionName"] for r in resp["Regions"]]
-
-            if self.conf["regions"]:
-                # filter out regions that are not available in the current partition
-                image_regions_configured_set = set(self.conf["regions"])
-                image_regions_all_set = set(image_regions_all)
-                self._image_regions = list(image_regions_configured_set.intersection(image_regions_all_set))
-                diff = image_regions_configured_set.difference(image_regions_all_set)
-                if diff:
-                    logger.warning(f"configured regions {diff} not available in the current partition. Ignoring those.")
-            else:
-                self._image_regions = image_regions_all
+            regions_configured = self.conf["regions"] if "regions" in self.conf else []
+            self._image_regions = _get_regions(self._s3.bucket_region, regions_configured)
             self._image_regions_cached = True
         return self._image_regions
 
@@ -358,14 +341,16 @@ def _sns_publish(self) -> None:
         """
         Publish SNS notifiations about newly available images to subscribers
         """
+
+        # Checking if the image(s) are registered or published before sending the notification
         for region in self.image_regions:
             ec2client_region: EC2Client = boto3.client("ec2", region_name=region)
             image_info: Optional[_ImageInfo] = self._get(ec2client_region)
-
             if not image_info:
                 logger.error(f"can not send SNS notification for {self.image_name} because no image found in {region}")
                 return
-            SNSNotification(self._ctx, self.image_name, region).publish()
+
+        SNSNotification(self._ctx, self.image_name, self._s3.bucket_region).publish()
 
     def cleanup(self) -> None:
         """

awspub/sns.py

@@ -11,8 +11,10 @@
 from mypy_boto3_sns.client import SNSClient
 from mypy_boto3_sts.client import STSClient
 
+from awspub.common import _get_regions
 from awspub.context import Context
 from awspub.exceptions import AWSAuthorizationException, AWSNotificationException
+from awspub.s3 import S3
 
 logger = logging.getLogger(__name__)
 
@@ -23,13 +25,13 @@ class SNSNotification(object):
     structuring rules for SNS notification JSON
     """
 
-    def __init__(self, context: Context, image_name: str, region_name: str):
+    def __init__(self, context: Context, image_name: str):
         """
         Construct a message and verify that it is valid
         """
         self._ctx: Context = context
         self._image_name: str = image_name
-        self._region_name: str = region_name
+        self._s3: S3 = S3(context)
 
     @property
     def conf(self) -> List[Dict[str, Any]]:
@@ -39,6 +41,20 @@ def conf(self) -> List[Dict[str, Any]]:
         return self._ctx.conf["images"][self._image_name]["sns"]
 
     def _get_topic_arn(self, topic_name: str) -> str:
+    def _sns_regions(self, topic_config: Dict[Any, Any]) -> List[str]:
+        """
+        Get the sns regions. Either configured in the sns configuration
+        or all available regions.
+        If a region is listed that is not available in the currently used partition,
+        that region will be ignored (eg. having us-east-1 configured but running in the aws-cn
+        partition doesn't include us-east-1 here).
+        """
+
+        regions_configured = topic_config["regions"] if "regions" in topic_config else []
+        sns_regions = _get_regions(self._s3.bucket_region, regions_configured)
+
+        return sns_regions
+
         """
         Calculate topic ARN based on partition, region, account and topic name
         :param topic_name: Name of topic

awspub/tests/test_common.py

@@ -1,6 +1,8 @@
+from unittest.mock import patch
+
 import pytest
 
-from awspub.common import _split_partition
+from awspub.common import _get_regions, _split_partition
 
 
 @pytest.mark.parametrize(
@@ -14,3 +16,19 @@
 )
 def test_common__split_partition(input, expected_output):
     assert _split_partition(input) == expected_output
+
+
+@pytest.mark.parametrize(
+    "regions_in_partition,configured_regions,expected_output",
+    [
+        (["region-1", "region-2"], ["region-1", "region-3"], ["region-1"]),
+        (["region-1", "region-2", "region-3"], ["region-4", "region-5"], []),
+        (["region-1", "region-2"], [], ["region-1", "region-2"]),
+    ],
+)
+def test_common__get_regions(regions_in_partition, configured_regions, expected_output):
+    with patch("boto3.client") as bclient_mock:
+        instance = bclient_mock.return_value
+        instance.describe_regions.return_value = {"Regions": [{"RegionName": r} for r in regions_in_partition]}
+
+        assert _get_regions("", configured_regions) == expected_output