Repository: getmoto/moto
Pull request #8361 (2024-12-03 16:42:46)

Athena: list_query_execution() not filtering by workgroup name
Added filtering ability to list_query_execution() using the `WorkGroup` param to better align the function to the documentation.

Changed the types for workgroup in `start_query_execution` as they did not reflect what was being passed in during the tests. 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/getmoto/moto/pull/8361?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 94.58%. Comparing base [(`42cbccf`)](https://app.codecov.io/gh/getmoto/moto/commit/42cbccf56885c45d5b0ccec0c72cf6cb58f1d840?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) to head [(`1b811b4`)](https://app.codecov.io/gh/getmoto/moto/commit/1b811b40b91bdad4718807a4f839812d05694998?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8361   +/-   ##
=======================================
  Coverage   94.58%   94.58%           
=======================================
  Files        1159     1159           
  Lines      101143   101146    +3     
=======================================
+ Hits        95669    95672    +3     
  Misses       5474     5474           
```

| [Flag](https://app.codecov.io/gh/getmoto/moto/pull/8361/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | Coverage Δ | |
|---|---|---|
| [servertests](https://app.codecov.io/gh/getmoto/moto/pull/8361/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `28.90% <20.00%> (-0.01%)` | :arrow_down: |
| [unittests](https://app.codecov.io/gh/getmoto/moto/pull/8361/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `94.56% <100.00%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/getmoto/moto/pull/8361?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).


Patch:
diff --git a/moto/athena/models.py b/moto/athena/models.py
index 988cd387a..c9cb498c1 100644
--- a/moto/athena/models.py
+++ b/moto/athena/models.py
@@ -107,7 +107,7 @@ class Execution(BaseModel):
         query: str,
         context: str,
         config: Dict[str, Any],
-        workgroup: WorkGroup,
+        workgroup: Optional[WorkGroup],
         execution_parameters: Optional[List[str]],
     ):
         self.id = str(mock_random.uuid4())
@@ -246,14 +246,14 @@ class AthenaBackend(BaseBackend):
         query: str,
         context: str,
         config: Dict[str, Any],
-        workgroup: WorkGroup,
+        workgroup: str,
         execution_parameters: Optional[List[str]],
     ) -> str:
         execution = Execution(
             query=query,
             context=context,
             config=config,
-            workgroup=workgroup,
+            workgroup=self.work_groups.get(workgroup),
             execution_parameters=execution_parameters,
         )
         self.executions[execution.id] = execution
@@ -271,7 +271,13 @@ class AthenaBackend(BaseBackend):
     def get_query_execution(self, exec_id: str) -> Execution:
         return self.executions[exec_id]
 
-    def list_query_executions(self) -> Dict[str, Execution]:
+    def list_query_executions(self, workgroup: Optional[str]) -> Dict[str, Execution]:
+        if workgroup is not None:
+            return {
+                exec_id: execution
+                for exec_id, execution in self.executions.items()
+                if execution.workgroup and execution.workgroup.name == workgroup
+            }
         return self.executions
 
     def get_query_results(self, exec_id: str) -> QueryResults:
diff --git a/moto/athena/responses.py b/moto/athena/responses.py
index 938f9a9ea..23b074160 100644
--- a/moto/athena/responses.py
+++ b/moto/athena/responses.py
@@ -104,7 +104,8 @@ class AthenaResponse(BaseResponse):
         return json.dumps(result.to_dict())
 
     def list_query_executions(self) -> str:
-        executions = self.athena_backend.list_query_executions()
+        workgroup = self._get_param("WorkGroup")
+        executions = self.athena_backend.list_query_executions(workgroup)
         return json.dumps({"QueryExecutionIds": [i for i in executions.keys()]})
 
     def stop_query_execution(self) -> str:
diff --git a/moto/memorydb/models.py b/moto/memorydb/models.py
index 37396b589..e6fd366c1 100644
--- a/moto/memorydb/models.py
+++ b/moto/memorydb/models.py
@@ -369,7 +369,10 @@ class MemoryDBBackend(BaseBackend):
         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 @@ class MemoryDBBackend(BaseBackend):
             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 @@ class MemoryDBBackend(BaseBackend):
             source=source,
         )
         self.snapshots[snapshot_name] = snapshot
+        if tags:
+            self.tag_resource(snapshot.arn, tags)
         return snapshot
 
     def describe_clusters(
@@ -561,16 +568,27 @@ class MemoryDBBackend(BaseBackend):
 
     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"]
 
@@ -578,8 +596,13 @@ class MemoryDBBackend(BaseBackend):
         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]


Test patch:
diff --git a/tests/test_athena/test_athena.py b/tests/test_athena/test_athena.py
index c48c08f6a..e539c42b1 100644
--- a/tests/test_athena/test_athena.py
+++ b/tests/test_athena/test_athena.py
@@ -438,6 +438,50 @@ def test_list_query_executions():
     assert executions["QueryExecutionIds"][0] == exec_id
 
 
+@mock_aws
+def test_list_query_executions_by_workgroup():
+    client = boto3.client("athena", region_name="us-east-1")
+
+    create_basic_workgroup(client=client, name="athena_workgroup")
+    create_basic_workgroup(client=client, name="athena_workgroup_1")
+
+    client.start_query_execution(
+        QueryString="query1",
+        QueryExecutionContext={"Database": "string"},
+        ResultConfiguration={"OutputLocation": "string"},
+        WorkGroup="athena_workgroup",
+    )
+    exec_result = client.start_query_execution(
+        QueryString="query1",
+        QueryExecutionContext={"Database": "string"},
+        ResultConfiguration={"OutputLocation": "string"},
+        WorkGroup="athena_workgroup_1",
+    )
+    exec_id = exec_result["QueryExecutionId"]
+
+    executions = client.list_query_executions(WorkGroup="athena_workgroup_1")
+    assert len(executions["QueryExecutionIds"]) == 1
+    assert executions["QueryExecutionIds"][0] == exec_id
+
+
+@mock_aws
+def test_list_query_executions_by_workgroup_when_none_match():
+    client = boto3.client("athena", region_name="us-east-1")
+
+    create_basic_workgroup(client=client, name="athena_workgroup")
+    create_basic_workgroup(client=client, name="athena_workgroup_1")
+
+    client.start_query_execution(
+        QueryString="query1",
+        QueryExecutionContext={"Database": "string"},
+        ResultConfiguration={"OutputLocation": "string"},
+        WorkGroup="athena_workgroup",
+    )
+
+    executions = client.list_query_executions(WorkGroup="athena_workgroup_1")
+    assert len(executions["QueryExecutionIds"]) == 0
+
+
 @mock_aws
 def test_list_named_queries():
     client = boto3.client("athena", region_name="us-east-1")
diff --git a/tests/test_memorydb/test_memorydb.py b/tests/test_memorydb/test_memorydb.py
index 8b4ff0c5c..6f20771cf 100644
--- a/tests/test_memorydb/test_memorydb.py
+++ b/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")