From 6848d55319ab4a815b6716716f4a5dfa515e6850 Mon Sep 17 00:00:00 2001
From: Darren Cohen <39422044+dargilco@users.noreply.github.com>
Date: Fri, 7 Feb 2025 16:09:26 -0800
Subject: [PATCH 1/2] First

---
 .../azure-ai-inference/azure/ai/inference/models/_patch.py   | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/sdk/ai/azure-ai-inference/azure/ai/inference/models/_patch.py b/sdk/ai/azure-ai-inference/azure/ai/inference/models/_patch.py
index 8744396de999..4b000f6b066c 100644
--- a/sdk/ai/azure-ai-inference/azure/ai/inference/models/_patch.py
+++ b/sdk/ai/azure-ai-inference/azure/ai/inference/models/_patch.py
@@ -6,7 +6,6 @@
 
 Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
 """
-import asyncio
 import base64
 import json
 import logging
@@ -503,8 +502,8 @@ async def _read_next_block_async(self) -> bool:
             return True
         return self._deserialize_and_add_to_queue(element)
 
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:  # type: ignore
-        asyncio.run(self.aclose())
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:  # type: ignore
+        await self.aclose()
 
     async def aclose(self) -> None:
         await self._response.close()

From 3f8b1657b17a26742d8822e064d7229ab89142af Mon Sep 17 00:00:00 2001
From: Darren Cohen <39422044+dargilco@users.noreply.github.com>
Date: Fri, 7 Feb 2025 22:30:03 -0800
Subject: [PATCH 2/2] More

---
 .../azure/ai/inference/models/_patch.py       |  6 ++++++
 .../tests/test_chat_completions_client.py     | 19 +++++++++----------
 .../test_chat_completions_client_async.py     | 19 +++++++++----------
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/sdk/ai/azure-ai-inference/azure/ai/inference/models/_patch.py b/sdk/ai/azure-ai-inference/azure/ai/inference/models/_patch.py
index 4b000f6b066c..d78ad267b36d 100644
--- a/sdk/ai/azure-ai-inference/azure/ai/inference/models/_patch.py
+++ b/sdk/ai/azure-ai-inference/azure/ai/inference/models/_patch.py
@@ -463,6 +463,9 @@ def _read_next_block(self) -> bool:
             return True
         return self._deserialize_and_add_to_queue(element)
 
+    def __enter__(self):
+        return self
+
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:  # type: ignore
         self.close()
 
@@ -502,6 +505,9 @@ async def _read_next_block_async(self) -> bool:
             return True
         return self._deserialize_and_add_to_queue(element)
 
+    async def __aenter__(self):
+        return self
+
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:  # type: ignore
         await self.aclose()
 
diff --git a/sdk/ai/azure-ai-inference/tests/test_chat_completions_client.py b/sdk/ai/azure-ai-inference/tests/test_chat_completions_client.py
index 269d775df75a..3626292153d2 100644
--- a/sdk/ai/azure-ai-inference/tests/test_chat_completions_client.py
+++ b/sdk/ai/azure-ai-inference/tests/test_chat_completions_client.py
@@ -389,16 +389,15 @@ def test_chat_completions_with_bytes_input(self, **kwargs):
     @ServicePreparerChatCompletions()
     @recorded_by_proxy
     def test_chat_completions_streaming(self, **kwargs):
-        client = self._create_chat_client(**kwargs)
-        response = client.complete(
-            stream=True,
-            messages=[
-                sdk.models.SystemMessage("You are a helpful assistant."),
-                sdk.models.UserMessage("Give me 3 good reasons why I should exercise every day."),
-            ],
-        )
-        self._validate_chat_completions_streaming_result(response)
-        client.close()
+        with self._create_chat_client(**kwargs) as client:
+            with client.complete(
+                stream=True,
+                messages=[
+                    sdk.models.SystemMessage("You are a helpful assistant."),
+                    sdk.models.UserMessage("Give me 3 good reasons why I should exercise every day."),
+                ],
+            ) as response:
+                self._validate_chat_completions_streaming_result(response)
 
     @ServicePreparerChatCompletions()
     @recorded_by_proxy
diff --git a/sdk/ai/azure-ai-inference/tests/test_chat_completions_client_async.py b/sdk/ai/azure-ai-inference/tests/test_chat_completions_client_async.py
index 33a173c2db47..dc13fc1eba4a 100644
--- a/sdk/ai/azure-ai-inference/tests/test_chat_completions_client_async.py
+++ b/sdk/ai/azure-ai-inference/tests/test_chat_completions_client_async.py
@@ -371,16 +371,15 @@ async def test_async_chat_completions_with_model_extras(self, **kwargs):
     @ServicePreparerChatCompletions()
     @recorded_by_proxy_async
     async def test_async_chat_completions_streaming(self, **kwargs):
-        client = self._create_async_chat_client(Sync=False, **kwargs)
-        response = await client.complete(
-            stream=True,
-            messages=[
-                sdk.models.SystemMessage(content="You are a helpful assistant."),
-                sdk.models.UserMessage(content="Give me 3 good reasons why I should exercise every day."),
-            ],
-        )
-        await self._validate_async_chat_completions_streaming_result(response)
-        await client.close()
+        async with self._create_async_chat_client(Sync=False, **kwargs) as client:
+            async with await client.complete(
+                stream=True,
+                messages=[
+                    sdk.models.SystemMessage(content="You are a helpful assistant."),
+                    sdk.models.UserMessage(content="Give me 3 good reasons why I should exercise every day."),
+                ],
+            ) as response:
+                await self._validate_async_chat_completions_streaming_result(response)
 
     @ServicePreparerChatCompletions()
     @recorded_by_proxy_async