From cac976754d0ebfa9ce933f6dcdda5d19827708a7 Mon Sep 17 00:00:00 2001 From: donfiguerres Date: Fri, 7 Oct 2022 17:40:29 +0800 Subject: [PATCH] Implement mediaconnect flow entitlement methods (#5536) --- IMPLEMENTATION_COVERAGE.md | 6 +- docs/docs/services/mediaconnect.rst | 6 +- moto/mediaconnect/models.py | 74 +++++++- moto/mediaconnect/responses.py | 35 ++++ moto/mediaconnect/urls.py | 2 + tests/test_mediaconnect/test_mediaconnect.py | 168 +++++++++++++++++++ 6 files changed, 284 insertions(+), 7 deletions(-) diff --git a/IMPLEMENTATION_COVERAGE.md b/IMPLEMENTATION_COVERAGE.md index 195d01577..cce81fbd3 100644 --- a/IMPLEMENTATION_COVERAGE.md +++ b/IMPLEMENTATION_COVERAGE.md @@ -4020,7 +4020,7 @@ - [X] describe_flow - [ ] describe_offering - [ ] describe_reservation -- [ ] grant_flow_entitlements +- [x] grant_flow_entitlements - [ ] list_entitlements - [X] list_flows - [ ] list_offerings @@ -4031,13 +4031,13 @@ - [X] remove_flow_output - [ ] remove_flow_source - [X] remove_flow_vpc_interface -- [ ] revoke_flow_entitlement +- [x] revoke_flow_entitlement - [X] start_flow - [X] stop_flow - [X] tag_resource - [ ] untag_resource - [ ] update_flow -- [ ] update_flow_entitlement +- [x] update_flow_entitlement - [ ] update_flow_media_stream - [ ] update_flow_output - [X] update_flow_source diff --git a/docs/docs/services/mediaconnect.rst b/docs/docs/services/mediaconnect.rst index 311121e84..f90c8d40b 100644 --- a/docs/docs/services/mediaconnect.rst +++ b/docs/docs/services/mediaconnect.rst @@ -34,7 +34,7 @@ mediaconnect - [X] describe_flow - [ ] describe_offering - [ ] describe_reservation -- [ ] grant_flow_entitlements +- [x] grant_flow_entitlements - [ ] list_entitlements - [X] list_flows - [ ] list_offerings @@ -45,13 +45,13 @@ mediaconnect - [X] remove_flow_output - [ ] remove_flow_source - [X] remove_flow_vpc_interface -- [ ] revoke_flow_entitlement +- [x] revoke_flow_entitlement - [X] start_flow - [X] stop_flow - [X] tag_resource - [ ] untag_resource - [ ] update_flow -- [ ] update_flow_entitlement +- [x] update_flow_entitlement - [ ] update_flow_media_stream - [ ] update_flow_output - [X] update_flow_source diff --git a/moto/mediaconnect/models.py b/moto/mediaconnect/models.py index d8e269a43..ffc4a243b 100644 --- a/moto/mediaconnect/models.py +++ b/moto/mediaconnect/models.py @@ -86,6 +86,14 @@ class MediaConnectBackend(BaseBackend): if not source.get("entitlementArn"): source["ingestIp"] = ingest_ip + def _add_entitlement_details(self, entitlement, entitlement_id): + if entitlement: + entitlement["entitlementArn"] = ( + f"arn:aws:mediaconnect:{self.region_name}" + f":{self.account_id}:entitlement:{entitlement_id}" + f":{entitlement['name']}" + ) + def _create_flow_add_details(self, flow): flow_id = random.uuid4().hex @@ -100,6 +108,10 @@ class MediaConnectBackend(BaseBackend): if output.get("protocol") in ["srt-listener", "zixi-pull"]: output["listenerAddress"] = f"{index}.0.0.0" + for _, entitlement in enumerate(flow.entitlements): + entitlement_id = random.uuid4().hex + self._add_entitlement_details(entitlement, entitlement_id) + def create_flow( self, availability_zone, @@ -306,7 +318,67 @@ class MediaConnectBackend(BaseBackend): source["whitelistCidr"] = whitelist_cidr return flow_arn, source - # add methods from here + def grant_flow_entitlements( + self, + flow_arn, + entitlements, + ): + if flow_arn not in self._flows: + raise NotFoundException( + message="flow with arn={} not found".format(flow_arn) + ) + flow = self._flows[flow_arn] + for entitlement in entitlements: + entitlement_id = random.uuid4().hex + name = entitlement["name"] + arn = f"arn:aws:mediaconnect:{self.region_name}:{self.account_id}:entitlement:{entitlement_id}:{name}" + entitlement["entitlementArn"] = arn + + flow.entitlements += entitlements + return flow_arn, entitlements + + def revoke_flow_entitlement(self, flow_arn, entitlement_arn): + if flow_arn not in self._flows: + raise NotFoundException( + message="flow with arn={} not found".format(flow_arn) + ) + flow = self._flows[flow_arn] + for entitlement in flow.entitlements: + if entitlement_arn == entitlement["entitlementArn"]: + flow.entitlements.remove(entitlement) + return flow_arn, entitlement_arn + raise NotFoundException( + message="entitlement with arn={} not found".format(entitlement_arn) + ) + + def update_flow_entitlement( + self, + flow_arn, + entitlement_arn, + description, + encryption, + entitlement_status, + name, + subscribers, + ): + if flow_arn not in self._flows: + raise NotFoundException( + message="flow with arn={} not found".format(flow_arn) + ) + flow = self._flows[flow_arn] + for entitlement in flow.entitlements: + if entitlement_arn == entitlement["entitlementArn"]: + entitlement["description"] = description + entitlement["encryption"] = encryption + entitlement["entitlementStatus"] = entitlement_status + entitlement["name"] = name + entitlement["subscribers"] = subscribers + return flow_arn, entitlement + raise NotFoundException( + message="entitlement with arn={} not found".format(entitlement_arn) + ) + + # add methods from here mediaconnect_backends = BackendDict(MediaConnectBackend, "mediaconnect") diff --git a/moto/mediaconnect/responses.py b/moto/mediaconnect/responses.py index 99b4fb034..7e1687f15 100644 --- a/moto/mediaconnect/responses.py +++ b/moto/mediaconnect/responses.py @@ -161,3 +161,38 @@ class MediaConnectResponse(BaseResponse): whitelist_cidr=whitelist_cidr, ) return json.dumps(dict(flow_arn=flow_arn, source=source)) + + def grant_flow_entitlements(self): + flow_arn = unquote(self._get_param("flowArn")) + entitlements = self._get_param("entitlements") + flow_arn, entitlements = self.mediaconnect_backend.grant_flow_entitlements( + flow_arn=flow_arn, entitlements=entitlements + ) + return json.dumps(dict(flow_arn=flow_arn, entitlements=entitlements)) + + def revoke_flow_entitlement(self): + flow_arn = unquote(self._get_param("flowArn")) + entitlement_arn = unquote(self._get_param("entitlementArn")) + flow_arn, entitlement_arn = self.mediaconnect_backend.revoke_flow_entitlement( + flow_arn=flow_arn, entitlement_arn=entitlement_arn + ) + return json.dumps(dict(flowArn=flow_arn, entitlementArn=entitlement_arn)) + + def update_flow_entitlement(self): + flow_arn = unquote(self._get_param("flowArn")) + entitlement_arn = unquote(self._get_param("entitlementArn")) + description = self._get_param("description") + encryption = self._get_param("encryption") + entitlement_status = self._get_param("entitlementStatus") + name = self._get_param("name") + subscribers = self._get_param("subscribers") + flow_arn, entitlement = self.mediaconnect_backend.update_flow_entitlement( + flow_arn=flow_arn, + entitlement_arn=entitlement_arn, + description=description, + encryption=encryption, + entitlement_status=entitlement_status, + name=name, + subscribers=subscribers, + ) + return json.dumps(dict(flowArn=flow_arn, entitlement=entitlement)) diff --git a/moto/mediaconnect/urls.py b/moto/mediaconnect/urls.py index cf9225ca0..1a9543ff5 100644 --- a/moto/mediaconnect/urls.py +++ b/moto/mediaconnect/urls.py @@ -17,6 +17,8 @@ url_paths = { "{0}/v1/flows/(?P[^/.]+)/source/(?P[^/.]+)": response.dispatch, "{0}/v1/flows/(?P[^/.]+)/outputs": response.dispatch, "{0}/v1/flows/(?P[^/.]+)/outputs/(?P[^/.]+)": response.dispatch, + "{0}/v1/flows/(?P[^/.]+)/entitlements": response.dispatch, + "{0}/v1/flows/(?P[^/.]+)/entitlements/(?P[^/.]+)": response.dispatch, "{0}/v1/flows/start/(?P[^/.]+)": response.dispatch, "{0}/v1/flows/stop/(?P[^/.]+)": response.dispatch, "{0}/tags/(?P[^/.]+)": response.dispatch, diff --git a/tests/test_mediaconnect/test_mediaconnect.py b/tests/test_mediaconnect/test_mediaconnect.py index 047dc829b..962245aab 100644 --- a/tests/test_mediaconnect/test_mediaconnect.py +++ b/tests/test_mediaconnect/test_mediaconnect.py @@ -471,3 +471,171 @@ def test_update_flow_source_succeeds(): FlowArn=flow_arn, SourceArn=source_arn, Description="new description" ) update_response["Source"]["Description"].should.equal("new description") + + +@mock_mediaconnect +def test_grant_flow_entitlements_fails(): + client = boto3.client("mediaconnect", region_name=region) + flow_arn = "unknown-flow" + + channel_config = _create_flow_config("test-Flow-1") + client.create_flow(**channel_config) + + with pytest.raises(ClientError) as err: + client.grant_flow_entitlements( + FlowArn=flow_arn, + Entitlements=[ + { + "DataTransferSubscriberFeePercent": 12, + "Description": "A new entitlement", + "Encryption": {"Algorithm": "aes256", "RoleArn": "some:role"}, + "EntitlementStatus": "ENABLED", + "Name": "Entitlement-B", + "Subscribers": [], + } + ], + ) + err = err.value.response["Error"] + err["Code"].should.equal("NotFoundException") + err["Message"].should.equal("flow with arn=unknown-flow not found") + + +@mock_mediaconnect +def test_grant_flow_entitlements_succeeds(): + client = boto3.client("mediaconnect", region_name=region) + channel_config = _create_flow_config("test-Flow-1") + + create_response = client.create_flow(**channel_config) + create_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + create_response["Flow"]["Status"].should.equal("STANDBY") + flow_arn = create_response["Flow"]["FlowArn"] + + describe_response = client.describe_flow(FlowArn=flow_arn) + describe_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + len(describe_response["Flow"]["Sources"]).should.equal(1) + + grant_response = client.grant_flow_entitlements( + FlowArn=flow_arn, + Entitlements=[ + { + "DataTransferSubscriberFeePercent": 12, + "Description": "A new entitlement", + "Encryption": {"Algorithm": "aes256", "RoleArn": "some:role"}, + "EntitlementStatus": "ENABLED", + "Name": "Entitlement-B", + "Subscribers": [], + }, + { + "DataTransferSubscriberFeePercent": 12, + "Description": "Another new entitlement", + "Encryption": {"Algorithm": "aes256", "RoleArn": "some:role"}, + "EntitlementStatus": "ENABLED", + "Name": "Entitlement-C", + "Subscribers": [], + }, + ], + ) + + entitlements = grant_response["Entitlements"] + len(entitlements).should.equal(2) + entitlement_names = [entitlement["Name"] for entitlement in entitlements] + entitlement_names.should.have("Entitlement-B") + entitlement_names.should.have("Entitlement-C") + + +@mock_mediaconnect +def test_revoke_flow_entitlement_fails(): + client = boto3.client("mediaconnect", region_name=region) + channel_config = _create_flow_config("test-Flow-1") + + create_response = client.create_flow(**channel_config) + create_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + create_response["Flow"]["Status"].should.equal("STANDBY") + flow_arn = create_response["Flow"]["FlowArn"] + + describe_response = client.describe_flow(FlowArn=flow_arn) + describe_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + len(describe_response["Flow"]["Entitlements"]).should.equal(1) + + with pytest.raises(ClientError) as err: + client.revoke_flow_entitlement( + FlowArn=flow_arn, EntitlementArn="some-other-arn" + ) + err = err.value.response["Error"] + err["Code"].should.equal("NotFoundException") + err["Message"].should.equal("entitlement with arn=some-other-arn not found") + + +@mock_mediaconnect +def test_revoke_flow_entitlement_succeeds(): + client = boto3.client("mediaconnect", region_name=region) + channel_config = _create_flow_config("test-Flow-1") + + create_response = client.create_flow(**channel_config) + create_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + create_response["Flow"]["Status"].should.equal("STANDBY") + flow_arn = create_response["Flow"]["FlowArn"] + + describe_response = client.describe_flow(FlowArn=flow_arn) + describe_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + entitlement_arn = describe_response["Flow"]["Entitlements"][0]["EntitlementArn"] + + revoke_response = client.revoke_flow_entitlement( + FlowArn=flow_arn, EntitlementArn=entitlement_arn + ) + revoke_response["FlowArn"].should.equal(flow_arn) + revoke_response["EntitlementArn"].should.equal(entitlement_arn) + + describe_response = client.describe_flow(FlowArn=flow_arn) + describe_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + len(describe_response["Flow"]["Entitlements"]).should.equal(0) + + +@mock_mediaconnect +def test_update_flow_entitlement_fails(): + client = boto3.client("mediaconnect", region_name=region) + channel_config = _create_flow_config("test-Flow-1") + + create_response = client.create_flow(**channel_config) + create_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + create_response["Flow"]["Status"].should.equal("STANDBY") + flow_arn = create_response["Flow"]["FlowArn"] + + describe_response = client.describe_flow(FlowArn=flow_arn) + describe_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + len(describe_response["Flow"]["Entitlements"]).should.equal(1) + + with pytest.raises(ClientError) as err: + client.update_flow_entitlement( + FlowArn=flow_arn, + EntitlementArn="some-other-arn", + Description="new description", + ) + err = err.value.response["Error"] + err["Code"].should.equal("NotFoundException") + err["Message"].should.equal("entitlement with arn=some-other-arn not found") + + +@mock_mediaconnect +def test_update_flow_entitlement_succeeds(): + client = boto3.client("mediaconnect", region_name=region) + channel_config = _create_flow_config("test-Flow-1") + + create_response = client.create_flow(**channel_config) + create_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + create_response["Flow"]["Status"].should.equal("STANDBY") + flow_arn = create_response["Flow"]["FlowArn"] + + describe_response = client.describe_flow(FlowArn=flow_arn) + describe_response["ResponseMetadata"]["HTTPStatusCode"].should.equal(200) + entitlement_arn = describe_response["Flow"]["Entitlements"][0]["EntitlementArn"] + + update_response = client.update_flow_entitlement( + FlowArn=flow_arn, + EntitlementArn=entitlement_arn, + Description="new description", + ) + update_response["FlowArn"].should.equal(flow_arn) + entitlement = update_response["Entitlement"] + entitlement["EntitlementArn"].should.equal(entitlement_arn) + entitlement["Description"].should.equal("new description")