Implement mediaconnect flow entitlement methods (#5536)

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

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

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")

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))

moto/mediaconnect/urls.py

@@ -17,6 +17,8 @@ url_paths = {
     "{0}/v1/flows/(?P<flowarn>[^/.]+)/source/(?P<sourcearn>[^/.]+)": response.dispatch,
     "{0}/v1/flows/(?P<flowarn>[^/.]+)/outputs": response.dispatch,
     "{0}/v1/flows/(?P<flowarn>[^/.]+)/outputs/(?P<outputarn>[^/.]+)": response.dispatch,
+    "{0}/v1/flows/(?P<flowarn>[^/.]+)/entitlements": response.dispatch,
+    "{0}/v1/flows/(?P<flowarn>[^/.]+)/entitlements/(?P<entitlementarn>[^/.]+)": response.dispatch,
     "{0}/v1/flows/start/(?P<flowarn>[^/.]+)": response.dispatch,
     "{0}/v1/flows/stop/(?P<flowarn>[^/.]+)": response.dispatch,
     "{0}/tags/(?P<resourcearn>[^/.]+)": response.dispatch,

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")