MemoryDB: adding tagging functionality to snapshots and subnets

moto/memorydb/models.py

@@ -369,7 +369,10 @@ def get_default_subnets(self) -> List[str]:
         return default_subnet_ids
 
     def _list_arns(self) -> List[str]:
-        return [cluster.arn for cluster in self.clusters.values()]
+        cluster_arns = [cluster.arn for cluster in self.clusters.values()]
+        snapshot_arns = [snapshot.arn for snapshot in self.snapshots.values()]
+        subnet_group_arns = [subnet.arn for subnet in self.subnet_groups.values()]
+        return cluster_arns + snapshot_arns + subnet_group_arns
 
     def create_cluster(
         self,
@@ -455,6 +458,8 @@ def create_subnet_group(
             tags,
         )
         self.subnet_groups[subnet_group_name] = subnet_group
+        if tags:
+            self.tag_resource(subnet_group.arn, tags)
         return subnet_group
 
     def create_snapshot(
@@ -483,6 +488,8 @@ def create_snapshot(
             source=source,
         )
         self.snapshots[snapshot_name] = snapshot
+        if tags:
+            self.tag_resource(snapshot.arn, tags)
         return snapshot
 
     def describe_clusters(
@@ -561,25 +568,41 @@ def describe_subnet_groups(
 
     def list_tags(self, resource_arn: str) -> List[Dict[str, str]]:
         if resource_arn not in self._list_arns():
-            cluster_name = resource_arn.split("/")[-1]
-            raise ClusterNotFoundFault(f"{cluster_name} is not present")
+            # Get the resource name from the resource_arn
+            resource_name = resource_arn.split("/")[-1]
+            if "subnetgroup" in resource_arn:
+                raise SubnetGroupNotFoundFault(f"{resource_name} is not present")
+            elif "snapshot" in resource_arn:
+                raise SnapshotNotFoundFault(f"{resource_name} is not present")
+            else:
+                raise ClusterNotFoundFault(f"{resource_name} is not present")
         return self.tagger.list_tags_for_resource(arn=resource_arn)["Tags"]
 
     def tag_resource(
         self, resource_arn: str, tags: List[Dict[str, str]]
     ) -> List[Dict[str, str]]:
         if resource_arn not in self._list_arns():
-            cluster_name = resource_arn.split("/")[-1]
-            raise ClusterNotFoundFault(f"{cluster_name} is not present")
+            resource_name = resource_arn.split("/")[-1]
+            if "subnetgroup" in resource_arn:
+                raise SubnetGroupNotFoundFault(f"{resource_name} is not present")
+            elif "snapshot" in resource_arn:
+                raise SnapshotNotFoundFault(f"{resource_name} is not present")
+            else:
+                raise ClusterNotFoundFault(f"{resource_name} is not present")
         self.tagger.tag_resource(resource_arn, tags)
         return self.tagger.list_tags_for_resource(arn=resource_arn)["Tags"]
 
     def untag_resource(
         self, resource_arn: str, tag_keys: List[str]
     ) -> List[Dict[str, str]]:
         if resource_arn not in self._list_arns():
-            cluster_name = resource_arn.split("/")[-1]
-            raise ClusterNotFoundFault(f"{cluster_name} is not present")
+            resource_name = resource_arn.split("/")[-1]
+            if "subnetgroup" in resource_arn:
+                raise SubnetGroupNotFoundFault(f"{resource_name} is not present")
+            elif "snapshot" in resource_arn:
+                raise SnapshotNotFoundFault(f"{resource_name} is not present")
+            else:
+                raise ClusterNotFoundFault(f"{resource_name} is not present")
         list_tags = self.list_tags(resource_arn=resource_arn)
         list_keys = [i["Key"] for i in list_tags]
         invalid_keys = [key for key in tag_keys if key not in list_keys]

tests/test_memorydb/test_memorydb.py

@@ -21,6 +21,9 @@ def create_subnet_group(client, region_name):
         SubnetGroupName="my_subnet_group",
         Description="This is my subnet group",
         SubnetIds=[subnet1.id, subnet2.id],
+        Tags=[
+            {"Key": "foo", "Value": "bar"},
+        ],
     )
     return subnet_group
 
@@ -395,6 +398,32 @@ def test_list_tags():
     assert "bar" in resp["TagList"][0]["Value"]
 
 
+@mock_aws
+def test_list_tags_snapshot():
+    client = boto3.client("memorydb", region_name="us-east-2")
+    client.create_cluster(
+        ClusterName="test-memory-db", NodeType="db.t4g.small", ACLName="open-access"
+    )
+    snapshot = client.create_snapshot(
+        ClusterName="test-memory-db",
+        SnapshotName="my-snapshot",
+        Tags=[
+            {"Key": "foo1", "Value": "bar1"},
+        ],
+    )
+    resp = client.list_tags(ResourceArn=snapshot["Snapshot"]["ARN"])
+    assert len(resp["TagList"]) == 1
+
+
+@mock_aws
+def test_list_tags_subnetgroups():
+    client = boto3.client("memorydb", region_name="us-east-2")
+    subnet_group = create_subnet_group(client, "us-east-2")
+    sg = subnet_group["SubnetGroup"]
+    resp = client.list_tags(ResourceArn=sg["ARN"])
+    assert len(resp["TagList"]) == 1
+
+
 @mock_aws
 def test_list_tags_invalid_cluster_fails():
     client = boto3.client("memorydb", region_name="us-east-2")
@@ -406,6 +435,28 @@ def test_list_tags_invalid_cluster_fails():
     assert err["Code"] == "ClusterNotFoundFault"
 
 
+@mock_aws
+def test_list_tags_invalid_snapshot_fails():
+    client = boto3.client("memorydb", region_name="us-east-2")
+    with pytest.raises(ClientError) as ex:
+        client.list_tags(
+            ResourceArn=f"arn:aws:memorydb:us-east-1:{ACCOUNT_ID}:snapshot/foobar",
+        )
+    err = ex.value.response["Error"]
+    assert err["Code"] == "SnapshotNotFoundFault"
+
+
+@mock_aws
+def test_list_tags_invalid_subnetgroup_fails():
+    client = boto3.client("memorydb", region_name="us-east-2")
+    with pytest.raises(ClientError) as ex:
+        client.list_tags(
+            ResourceArn=f"arn:aws:memorydb:us-east-1:{ACCOUNT_ID}:subnetgroup/foobar",
+        )
+    err = ex.value.response["Error"]
+    assert err["Code"] == "SubnetGroupNotFoundFault"
+
+
 @mock_aws
 def test_tag_resource():
     client = boto3.client("memorydb", region_name="us-east-2")
@@ -441,6 +492,30 @@ def test_tag_resource_invalid_cluster_fails():
     assert err["Code"] == "ClusterNotFoundFault"
 
 
+@mock_aws
+def test_tag_resource_invalid_snapshot_fails():
+    client = boto3.client("memorydb", region_name="us-east-2")
+    with pytest.raises(ClientError) as ex:
+        client.tag_resource(
+            ResourceArn=f"arn:aws:memorydb:us-east-1:{ACCOUNT_ID}:snapshot/foobar",
+            Tags=[{"Key": "key2", "Value": "value2"}],
+        )
+    err = ex.value.response["Error"]
+    assert err["Code"] == "SnapshotNotFoundFault"
+
+
+@mock_aws
+def test_tag_resource_invalid_subnetgroup_fails():
+    client = boto3.client("memorydb", region_name="us-east-2")
+    with pytest.raises(ClientError) as ex:
+        client.tag_resource(
+            ResourceArn=f"arn:aws:memorydb:us-east-1:{ACCOUNT_ID}:subnetgroup/foobar",
+            Tags=[{"Key": "key2", "Value": "value2"}],
+        )
+    err = ex.value.response["Error"]
+    assert err["Code"] == "SubnetGroupNotFoundFault"
+
+
 @mock_aws
 def test_untag_resource():
     client = boto3.client("memorydb", region_name="us-east-2")
@@ -474,6 +549,30 @@ def test_untag_resource_invalid_cluster_fails():
     assert err["Code"] == "ClusterNotFoundFault"
 
 
+@mock_aws
+def test_untag_resource_invalid_snapshot_fails():
+    client = boto3.client("memorydb", region_name="us-east-2")
+    with pytest.raises(ClientError) as ex:
+        client.untag_resource(
+            ResourceArn=f"arn:aws:memorydb:us-east-1:{ACCOUNT_ID}:snapshot/foobar",
+            TagKeys=["key1"],
+        )
+    err = ex.value.response["Error"]
+    assert err["Code"] == "SnapshotNotFoundFault"
+
+
+@mock_aws
+def test_untag_resource_invalid_subnetgroup_fails():
+    client = boto3.client("memorydb", region_name="us-east-2")
+    with pytest.raises(ClientError) as ex:
+        client.untag_resource(
+            ResourceArn=f"arn:aws:memorydb:us-east-1:{ACCOUNT_ID}:subnetgroup/foobar",
+            TagKeys=["key1"],
+        )
+    err = ex.value.response["Error"]
+    assert err["Code"] == "SubnetGroupNotFoundFault"
+
+
 @mock_aws
 def test_untag_resource_invalid_keys_fails():
     client = boto3.client("memorydb", region_name="us-east-2")