From 8b6dfe6e44aa051b0b22ab485f35f6080a1f45b7 Mon Sep 17 00:00:00 2001
From: Rakshith <rakshith.bhyravabhotla@gmail.com>
Date: Mon, 25 Jan 2021 01:25:08 -0800
Subject: [PATCH] Missing enter, exit and close

---
 .../azure/eventgrid/_publisher_client.py          | 15 +++++++++++++++
 .../eventgrid/aio/_publisher_client_async.py      | 12 ++++++++++++
 .../tests/test_eg_publisher_client_async.py       | 14 ++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/sdk/eventgrid/azure-eventgrid/azure/eventgrid/_publisher_client.py b/sdk/eventgrid/azure-eventgrid/azure/eventgrid/_publisher_client.py
index ae06329cbc48..4b197d171bc7 100644
--- a/sdk/eventgrid/azure-eventgrid/azure/eventgrid/_publisher_client.py
+++ b/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
diff --git a/sdk/eventgrid/azure-eventgrid/azure/eventgrid/aio/_publisher_client_async.py b/sdk/eventgrid/azure-eventgrid/azure/eventgrid/aio/_publisher_client_async.py
index cefa3f605d2c..a687c3916860 100644
--- a/sdk/eventgrid/azure-eventgrid/azure/eventgrid/aio/_publisher_client_async.py
+++ b/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__()
diff --git a/sdk/eventgrid/azure-eventgrid/tests/test_eg_publisher_client_async.py b/sdk/eventgrid/azure-eventgrid/tests/test_eg_publisher_client_async.py
index fd801f452b16..d8fdf2dcc659 100644
--- a/sdk/eventgrid/azure-eventgrid/tests/test_eg_publisher_client_async.py
+++ b/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)
\ No newline at end of file