Missing enter, exit and close (#16321)

Author: Rakshith Bhyravabhotla

sdk/eventgrid/azure-eventgrid/azure/eventgrid/_publisher_client.py

@@ -134,3 +134,18 @@ def send_events(self, events, **kwargs):
             self._client.publish_custom_event_events(self._endpoint, cast(List, serialized_events), **kwargs)
         else:
             raise ValueError("Event schema is not correct.")
+
+    def close(self):
+        # type: () -> None
+        """Close the :class:`~azure.eventgrid.EventGridPublisherClient` session.
+        """
+        return self._client.close()
+
+    def __enter__(self):
+        # type: () -> EventGridPublisherClient
+        self._client.__enter__()  # pylint:disable=no-member
+        return self
+
+    def __exit__(self, *args):
+        # type: (*Any) -> None
+        self._client.__exit__(*args)  # pylint:disable=no-member

sdk/eventgrid/azure-eventgrid/azure/eventgrid/aio/_publisher_client_async.py

@@ -146,3 +146,15 @@ async def send_events(
                 )
         else:
             raise ValueError("Event schema is not correct.")
+
+    async def __aenter__(self) -> "EventGridPublisherClient":
+        await self._client.__aenter__()
+        return self
+
+    async def __aexit__(self, *args: "Any") -> None:
+        await self._client.__aexit__(*args)
+
+    async def close(self) -> None:
+        """Close the :class:`~azure.eventgrid.aio.EventGridPublisherClient` session.
+        """
+        await self._client.__aexit__()

sdk/eventgrid/azure-eventgrid/tests/test_eg_publisher_client_async.py

@@ -274,3 +274,17 @@ async def test_send_custom_schema_event_as_list(self, resource_group, eventgrid_
                     }
                 )
         await client.send_events([custom_event1, custom_event2])
+
+    @CachedResourceGroupPreparer(name_prefix='eventgridtest')
+    @CachedEventGridTopicPreparer(name_prefix='cloudeventgridtest')
+    @pytest.mark.asyncio
+    async def test_send_and_close_async_session(self, resource_group, eventgrid_topic, eventgrid_topic_primary_key, eventgrid_topic_endpoint):
+        akc_credential = AzureKeyCredential(eventgrid_topic_primary_key)
+        client = EventGridPublisherClient(eventgrid_topic_endpoint, akc_credential)
+        async with client: # this throws if client can't close
+            cloud_event = CloudEvent(
+                    source = "http://samplesource.dev",
+                    data = "cloudevent",
+                    type="Sample.Cloud.Event"
+                    )
+            await client.send_events(cloud_event)