Update async wrapper

Author: Rena Chen

sdk/batch/azure-batch/azure/batch/_operations/_patch.py

@@ -7,7 +7,7 @@
 Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
 """
 import datetime
-from typing import Any, List, Optional
+from typing import Any, List, Optional, Iterator
 import collections
 import logging
 import threading
@@ -134,7 +134,7 @@ def get_node_file(
         if_unmodified_since: Optional[datetime.datetime] = None,
         ocp_range: Optional[str] = None,
         **kwargs: Any
-    ) -> bytes:
+    ) -> Iterator[bytes]:
         """Returns the content of the specified Compute Node file.
 
         :param pool_id: The ID of the Pool that contains the Compute Node. Required.
@@ -321,7 +321,7 @@ def get_task_file(
         if_unmodified_since: Optional[datetime.datetime] = None,
         ocp_range: Optional[str] = None,
         **kwargs: Any
-    ) -> bytes:
+    ) -> Iterator[bytes]:
         """Returns the content of the specified Task file.
 
         :param job_id: The ID of the Job that contains the Task. Required.

sdk/batch/azure-batch/azure/batch/aio/_operations/_patch.py

@@ -10,7 +10,7 @@
 import datetime
 import collections
 import logging
-from typing import Any, List, Optional
+from typing import Any, AsyncIterator, List, Optional
 
 from azure.batch import models as _models
 from azure.core import MatchConditions
@@ -125,7 +125,7 @@ async def get_node_file(
         if_unmodified_since: Optional[datetime.datetime] = None,
         ocp_range: Optional[str] = None,
         **kwargs: Any
-    ) -> bytes:
+    ) -> AsyncIterator[bytes]:
         """Returns the content of the specified Compute Node file.
 
         :param pool_id: The ID of the Pool that contains the Compute Node. Required.
@@ -312,7 +312,7 @@ async def get_task_file(
         if_unmodified_since: Optional[datetime.datetime] = None,
         ocp_range: Optional[str] = None,
         **kwargs: Any
-    ) -> bytes:
+    ) -> AsyncIterator[bytes]:
         """Returns the content of the specified Task file.
 
         :param job_id: The ID of the Job that contains the Task. Required.

sdk/batch/azure-batch/tests/async_wrapper.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import AsyncIterable, AsyncIterator, Any, cast, Coroutine, Iterable, Iterator, TypeVar, Union
+from typing import Any, AsyncIterator, AsyncIterable, cast, Coroutine, Iterable, Iterator, TypeVar, Union
 
 T = TypeVar("T")
 

sdk/batch/azure-batch/tests/test_batch.py

@@ -21,7 +21,7 @@
 from typing import Iterable, Union
 
 from batch_preparers import AccountPreparer, PoolPreparer, JobPreparer
-from async_wrapper import wrap_list_result, wrap_result
+from async_wrapper import wrap_list_result, wrap_result, wrap_file_result
 from decorators import recorded_by_proxy_async, client_setup
 
 from devtools_testutils import (
@@ -751,6 +751,7 @@ async def test_batch_files(self, client: BatchClient, **kwargs):
             nodes = list(await wrap_list_result(client.list_nodes(batch_pool.name)))
         assert len(nodes) == 1
         node = nodes[0].id
+        assert node is not None
         task_id = "test_task"
         task_param = models.BatchTaskCreateContent(id=task_id, command_line='cmd /c "echo hello world"')
         response = await wrap_result(client.create_task(batch_job.id, task_param))
@@ -773,7 +774,7 @@ async def test_batch_files(self, client: BatchClient, **kwargs):
         # Test Get File from Batch Node
         file_length = 0
         with io.BytesIO() as file_handle:
-            response = await wrap_result(client.get_node_file(batch_pool.name, node, only_files[0].name))
+            response = await wrap_file_result(client.get_node_file(batch_pool.name, node, only_files[0].name))
             for data in response:
                 file_length += 1
         assert file_length == int(props.headers["Content-Length"])
@@ -797,7 +798,7 @@ async def test_batch_files(self, client: BatchClient, **kwargs):
         # Test Get File from Task
         file_length = 0
         with io.BytesIO() as file_handle:
-            response = await wrap_result(client.get_task_file(batch_job.id, task_id, only_files[0].name))
+            response = await wrap_file_result(client.get_task_file(batch_job.id, task_id, only_files[0].name))
             for data in response:
                 file_length += len(data)
         assert file_length == int(props.headers["Content-Length"])