clean up to satisfy mypy

tests/asyncio/gapic/test_method_async.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 
 import datetime
-from unittest import mock
-
 try:
-    from unittest.mock import AsyncMock
-except ImportError:
-    from mock import AsyncMock
+    from unittest import mock
+    from unittest.mock import AsyncMock  # pragma: NO COVER
+except ImportError:  # pragma: NO COVER
+    import mock  # type: ignore
 import pytest
 
 try:
@@ -260,7 +259,7 @@ async def test_wrap_method_with_overriding_timeout_as_a_number():
 
 @pytest.mark.asyncio
 async def test_wrap_method_without_wrap_errors():
-    fake_call = AsyncMock()
+    fake_call = mock.AsyncMock()
 
     wrapped_method = gapic_v1.method_async.wrap_method(fake_call, kind="rest")
     with mock.patch("google.api_core.grpc_helpers_async.wrap_errors") as method:

tests/asyncio/retry/test_retry_unary_async.py

@@ -14,12 +14,11 @@
 
 import datetime
 import re
-from unittest import mock
-
 try:
-    from unittest.mock import AsyncMock
-except ImportError:
-    from mock import AsyncMock
+    from unittest import mock
+    from unittest.mock import AsyncMock  # pragma: NO COVER
+except ImportError:  # pragma: NO COVER
+    import mock  # type: ignore
 import pytest
 
 from google.api_core import exceptions
@@ -172,7 +171,7 @@ def if_exception_type(exc):
     @pytest.mark.asyncio
     async def test___call___and_execute_success(self, sleep):
         retry_ = retry_async.AsyncRetry()
-        target = AsyncMock(spec=["__call__"], return_value=42)
+        target = mock.AsyncMock(spec=["__call__"], return_value=42)
         # __name__ is needed by functools.partial.
         target.__name__ = "target"
 
@@ -194,7 +193,7 @@ async def test___call___and_execute_retry(self, sleep, uniform):
             predicate=retry_async.if_exception_type(ValueError)
         )
 
-        target = AsyncMock(spec=["__call__"], side_effect=[ValueError(), 42])
+        target = mock.AsyncMock(spec=["__call__"], side_effect=[ValueError(), 42])
         # __name__ is needed by functools.partial.
         target.__name__ = "target"
 
@@ -224,7 +223,7 @@ async def test___call___and_execute_retry_hitting_timeout(self, sleep, uniform):
 
         monotonic_patcher = mock.patch("time.monotonic", return_value=0)
 
-        target = AsyncMock(spec=["__call__"], side_effect=[ValueError()] * 10)
+        target = mock.AsyncMock(spec=["__call__"], side_effect=[ValueError()] * 10)
         # __name__ is needed by functools.partial.
         target.__name__ = "target"
 
@@ -274,7 +273,7 @@ async def test___init___without_retry_executed(self, sleep):
         # check the proper creation of the class
         assert retry_._on_error is _some_function
 
-        target = AsyncMock(spec=["__call__"], side_effect=[42])
+        target = mock.AsyncMock(spec=["__call__"], side_effect=[42])
         # __name__ is needed by functools.partial.
         target.__name__ = "target"
 
@@ -299,7 +298,7 @@ async def test___init___when_retry_is_executed(self, sleep, uniform):
         # check the proper creation of the class
         assert retry_._on_error is _some_function
 
-        target = AsyncMock(
+        target = mock.AsyncMock(
             spec=["__call__"], side_effect=[ValueError(), ValueError(), 42]
         )
         # __name__ is needed by functools.partial.

tests/asyncio/test_grpc_helpers_async.py

@@ -12,12 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest import mock
-
 try:
-    from unittest.mock import AsyncMock
-except ImportError:
-    from mock import AsyncMock
+    from unittest import mock
+    from unittest.mock import AsyncMock  # pragma: NO COVER
+except ImportError:  # pragma: NO COVER
+    import mock  # type: ignore
 import pytest  # noqa: I202
 
 try:
@@ -54,7 +53,7 @@ def trailing_metadata(self):
 @pytest.mark.asyncio
 async def test_wrap_unary_errors():
     grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT)
-    callable_ = AsyncMock(spec=["__call__"], side_effect=grpc_error)
+    callable_ = mock.AsyncMock(spec=["__call__"], side_effect=grpc_error)
 
     wrapped_callable = grpc_helpers_async._wrap_unary_errors(callable_)
 
@@ -174,7 +173,7 @@ async def test_wrap_stream_errors_stream_stream():
 async def test_wrap_stream_errors_raised():
     grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT)
     mock_call = mock.Mock(aio.StreamStreamCall, autospec=True)
-    mock_call.wait_for_connection = AsyncMock(side_effect=[grpc_error])
+    mock_call.wait_for_connection = mock.AsyncMock(side_effect=[grpc_error])
     multicallable = mock.Mock(return_value=mock_call)
 
     wrapped_callable = grpc_helpers_async._wrap_stream_errors(
@@ -191,7 +190,7 @@ async def test_wrap_stream_errors_read():
     grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT)
 
     mock_call = mock.Mock(aio.StreamStreamCall, autospec=True)
-    mock_call.read = AsyncMock(side_effect=grpc_error)
+    mock_call.read = mock.AsyncMock(side_effect=grpc_error)
     multicallable = mock.Mock(return_value=mock_call)
 
     wrapped_callable = grpc_helpers_async._wrap_stream_errors(
@@ -213,7 +212,7 @@ async def test_wrap_stream_errors_aiter():
 
     mock_call = mock.Mock(aio.StreamStreamCall, autospec=True)
     mocked_aiter = mock.Mock(spec=["__anext__"])
-    mocked_aiter.__anext__ = AsyncMock(side_effect=[mock.sentinel.response, grpc_error])
+    mocked_aiter.__anext__ = mock.AsyncMock(side_effect=[mock.sentinel.response, grpc_error])
     mock_call.__aiter__ = mock.Mock(return_value=mocked_aiter)
     multicallable = mock.Mock(return_value=mock_call)
 
@@ -234,7 +233,7 @@ async def test_wrap_stream_errors_aiter_non_rpc_error():
 
     mock_call = mock.Mock(aio.StreamStreamCall, autospec=True)
     mocked_aiter = mock.Mock(spec=["__anext__"])
-    mocked_aiter.__anext__ = AsyncMock(
+    mocked_aiter.__anext__ = mock.AsyncMock(
         side_effect=[mock.sentinel.response, non_grpc_error]
     )
     mock_call.__aiter__ = mock.Mock(return_value=mocked_aiter)
@@ -269,8 +268,8 @@ async def test_wrap_stream_errors_write():
     grpc_error = RpcErrorImpl(grpc.StatusCode.INVALID_ARGUMENT)
 
     mock_call = mock.Mock(aio.StreamStreamCall, autospec=True)
-    mock_call.write = AsyncMock(side_effect=[None, grpc_error])
-    mock_call.done_writing = AsyncMock(side_effect=[None, grpc_error])
+    mock_call.write = mock.AsyncMock(side_effect=[None, grpc_error])
+    mock_call.done_writing = mock.AsyncMock(side_effect=[None, grpc_error])
     multicallable = mock.Mock(return_value=mock_call)
 
     wrapped_callable = grpc_helpers_async._wrap_stream_errors(

tests/asyncio/test_operation_async.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 
 
-from unittest import mock
-
-try:
-    from unittest.mock import AsyncMock
-except ImportError:
-    from mock import AsyncMock
 import pytest
+try:
+    from unittest import mock
+    from unittest.mock import AsyncMock  # pragma: NO COVER
+except ImportError:  # pragma: NO COVER
+    import mock  # type: ignore
 
 try:
     import grpc  # noqa: F401
@@ -59,9 +58,9 @@ def make_operation_future(client_operations_responses=None):
     if client_operations_responses is None:
         client_operations_responses = [make_operation_proto()]
 
-    refresh = AsyncMock(spec=["__call__"], side_effect=client_operations_responses)
+    refresh = mock.AsyncMock(spec=["__call__"], side_effect=client_operations_responses)
     refresh.responses = client_operations_responses
-    cancel = AsyncMock(spec=["__call__"])
+    cancel = mock.AsyncMock(spec=["__call__"])
     operation_future = operation_async.AsyncOperation(
         client_operations_responses[0],
         refresh,

tests/asyncio/test_page_iterator_async.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 
 import inspect
-from unittest import mock
-
 try:
-    from unittest.mock import AsyncMock
-except ImportError:
-    from mock import AsyncMock
+    from unittest import mock
+    from unittest.mock import AsyncMock  # pragma: NO COVER
+except ImportError:  # pragma: NO COVER
+    import mock  # type: ignore
 import pytest
 
 from google.api_core import page_iterator_async
@@ -62,7 +61,7 @@ async def test_anext(self):
         )
 
         async_iterator = PageAsyncIteratorImpl(None, None)
-        async_iterator._next_page = AsyncMock(side_effect=[page_1, page_2, None])
+        async_iterator._next_page = mock.AsyncMock(side_effect=[page_1, page_2, None])
 
         # Consume items and check the state of the async_iterator.
         assert async_iterator.num_results == 0
@@ -102,7 +101,7 @@ async def test__page_aiter_increment(self):
         page = page_iterator_async.Page(
             iterator, ("item",), page_iterator_async._item_to_value_identity
         )
-        iterator._next_page = AsyncMock(side_effect=[page, None])
+        iterator._next_page = mock.AsyncMock(side_effect=[page, None])
 
         assert iterator.num_results == 0
 
@@ -140,7 +139,7 @@ async def test__items_aiter(self):
         )
 
         iterator = PageAsyncIteratorImpl(None, None)
-        iterator._next_page = AsyncMock(side_effect=[page1, page2, None])
+        iterator._next_page = mock.AsyncMock(side_effect=[page1, page2, None])
 
         items_aiter = iterator._items_aiter()
 
@@ -163,7 +162,7 @@ async def test__items_aiter(self):
     @pytest.mark.asyncio
     async def test___aiter__(self):
         async_iterator = PageAsyncIteratorImpl(None, None)
-        async_iterator._next_page = AsyncMock(side_effect=[(1, 2), (3,), None])
+        async_iterator._next_page = mock.AsyncMock(side_effect=[(1, 2), (3,), None])
 
         assert not async_iterator._started
 
@@ -252,7 +251,7 @@ async def test_iterate(self):
         response1 = mock.Mock(items=["a", "b"], next_page_token="1")
         response2 = mock.Mock(items=["c"], next_page_token="2")
         response3 = mock.Mock(items=["d"], next_page_token="")
-        method = AsyncMock(side_effect=[response1, response2, response3])
+        method = mock.AsyncMock(side_effect=[response1, response2, response3])
         iterator = page_iterator_async.AsyncGRPCIterator(
             mock.sentinel.client, method, request, "items"
         )
@@ -275,7 +274,7 @@ async def test_iterate_with_max_results(self):
         response1 = mock.Mock(items=["a", "b"], next_page_token="1")
         response2 = mock.Mock(items=["c"], next_page_token="2")
         response3 = mock.Mock(items=["d"], next_page_token="")
-        method = AsyncMock(side_effect=[response1, response2, response3])
+        method = mock.AsyncMock(side_effect=[response1, response2, response3])
         iterator = page_iterator_async.AsyncGRPCIterator(
             mock.sentinel.client, method, request, "items", max_results=3
         )

tests/asyncio/test_rest_streaming_async.py

@@ -20,12 +20,12 @@
 import random
 import time
 from typing import List, AsyncIterator
-from unittest import mock
-
 try:
-    from unittest.mock import AsyncMock
-except ImportError:
-    from mock import AsyncMock
+    from unittest import mock
+    from unittest.mock import AsyncMock  # pragma: NO COVER
+except ImportError:  # pragma: NO COVER
+    import mock  # type: ignore
+
 import pytest  # noqa: I202
 
 import proto
@@ -307,7 +307,7 @@ async def test_next_not_array(response_type):
 @pytest.mark.parametrize("response_type", [EchoResponse, httpbody_pb2.HttpBody])
 async def test_cancel(response_type):
     with mock.patch.object(
-        ResponseMock, "close", new_callable=AsyncMock
+        ResponseMock, "close", new_callable=mock.AsyncMock
     ) as mock_method:
         resp = ResponseMock(responses=[], response_cls=response_type)
         itr = rest_streaming_async.AsyncResponseIterator(resp, response_type)
@@ -319,7 +319,7 @@ async def test_cancel(response_type):
 @pytest.mark.parametrize("response_type", [EchoResponse, httpbody_pb2.HttpBody])
 async def test_iterator_as_context_manager(response_type):
     with mock.patch.object(
-        ResponseMock, "close", new_callable=AsyncMock
+        ResponseMock, "close", new_callable=mock.AsyncMock
     ) as mock_method:
         resp = ResponseMock(responses=[], response_cls=response_type)
         async with rest_streaming_async.AsyncResponseIterator(resp, response_type):

tests/unit/retry/test_retry_unary.py

@@ -13,14 +13,13 @@
 # limitations under the License.
 
 import datetime
+import pytest
 import re
-from unittest import mock
-
 try:
-    from unittest.mock import AsyncMock
-except ImportError:
-    from mock import AsyncMock
-import pytest
+    from unittest import mock
+    from unittest.mock import AsyncMock  # pragma: NO COVER
+except ImportError:  # pragma: NO COVER
+    import mock  # type: ignore
 
 from google.api_core import exceptions
 from google.api_core import retry
@@ -106,7 +105,7 @@ def test_retry_target_non_retryable_error(utcnow, sleep):
 @pytest.mark.asyncio
 async def test_retry_target_warning_for_retry(utcnow, sleep):
     predicate = retry.if_exception_type(ValueError)
-    target = AsyncMock(spec=["__call__"])
+    target = mock.AsyncMock(spec=["__call__"])
 
     with pytest.warns(Warning) as exc_info:
         # Note: predicate is just a filler and doesn't affect the test