add async system tests for resumable uploads

Author: parthea

packages/gapic-generator/tests/system/conftest.py

@@ -44,6 +44,7 @@
 except ImportError:
     HAS_RESUMABLE_UPLOAD_CLIENT = False
 
+HAS_ASYNC_REST_RESUMABLE_UPLOAD_TRANSPORT = False
 if os.environ.get("GAPIC_PYTHON_ASYNC", "true") == "true":
     from grpc.experimental import aio
     import asyncio
@@ -67,6 +68,16 @@
         HAS_ASYNC_REST_IDENTITY_TRANSPORT = True
     except:
         HAS_ASYNC_REST_IDENTITY_TRANSPORT = False
+    try:
+        from google.showcase import ResumableUploadServiceAsyncClient
+        from google.showcase_v1beta1.services.resumable_upload_service.transports import (
+            AsyncResumableUploadServiceRestTransport,
+            AsyncResumableUploadServiceRestInterceptor,
+        )
+
+        HAS_ASYNC_REST_RESUMABLE_UPLOAD_TRANSPORT = True
+    except:
+        HAS_ASYNC_REST_RESUMABLE_UPLOAD_TRANSPORT = False
 
     _GRPC_VERSION = grpc.__version__
 
@@ -344,6 +355,19 @@ async def post_expand_with_metadata(self, request, metadata):
             return request, metadata
 
 
+if HAS_ASYNC_REST_RESUMABLE_UPLOAD_TRANSPORT:
+
+    class ResumableUploadMetadataClientRestAsyncInterceptor(
+        AsyncResumableUploadServiceRestInterceptor
+    ):
+        request_metadata: Sequence[Tuple[str, str]] = []
+        response_metadata: Sequence[Tuple[str, str]] = []
+
+        async def pre_upload_media(self, request, metadata):
+            self.request_metadata = metadata
+            return request, metadata
+
+
 class EchoMetadataClientGrpcInterceptor(
     grpc.UnaryUnaryClientInterceptor,
     grpc.UnaryStreamClientInterceptor,
@@ -576,6 +600,27 @@ def intercepted_resumable_upload_rest(use_mtls, use_tls):
     return ResumableUploadServiceClient(transport=transport), interceptor
 
 
+@pytest.fixture
+def intercepted_resumable_upload_rest_async():
+    if not HAS_ASYNC_REST_RESUMABLE_UPLOAD_TRANSPORT:
+        pytest.skip("Skipping test with async rest.")
+
+    transport_name = "rest_asyncio"
+    transport_cls = ResumableUploadServiceAsyncClient.get_transport_class(
+        transport_name
+    )
+    interceptor = ResumableUploadMetadataClientRestAsyncInterceptor()
+
+    transport = transport_cls(
+        credentials=async_anonymous_credentials(),
+        host="localhost:7469",
+        url_scheme="http",
+        interceptor=interceptor,
+    )
+
+    return ResumableUploadServiceAsyncClient(transport=transport), interceptor
+
+
 try:
     from google.api_core.resumable_transfer import (
         ResumableUploadConfig,

packages/gapic-generator/tests/system/test_resumable_upload_progress.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import io
+import os
+import pytest
 
 from google.api_core.resumable_transfer import (
     DEFAULT_CHUNK_SIZE,
@@ -274,3 +276,59 @@ def test_client_upload_media_passes_request_body(intercepted_resumable_upload_re
     assert response.name == "client_upload_media.txt"
     assert response.size == len(payload)
 
+
+if os.environ.get("GAPIC_PYTHON_ASYNC", "true") == "true":
+
+    @pytest.mark.asyncio
+    async def test_async_client_upload_media_end_to_end(
+        intercepted_resumable_upload_rest_async,
+    ):
+        """Verify end-to-end async resumable upload via ResumableUploadServiceAsyncClient."""
+        client, _ = intercepted_resumable_upload_rest_async
+        payload = b"0123456789" * 100
+        stream = io.BytesIO(payload)
+
+        session = await client.upload_media(
+            request=UploadMediaRequest(name="async_client_upload_media.txt"),
+            config=ResumableUploadConfig(chunk_size=256),
+        )
+        response = await session.upload(stream=stream)
+
+        assert isinstance(response, UploadMediaResponse)
+        assert response.name == "async_client_upload_media.txt"
+        assert response.size == len(payload)
+
+    @pytest.mark.asyncio
+    async def test_async_client_upload_media_progress_tracking(
+        intercepted_resumable_upload_rest_async,
+    ):
+        """Verify async progress iteration over ResumableUploadServiceAsyncClient.upload_media."""
+        client, _ = intercepted_resumable_upload_rest_async
+        chunk_size = 262_144  # 256 KiB
+        total_size = 600_000
+        payload = b"A" * total_size
+        stream = io.BytesIO(payload)
+
+        session = await client.upload_media(
+            request=UploadMediaRequest(name="async_multi_chunk.bin"),
+            config=ResumableUploadConfig(chunk_size=chunk_size),
+        )
+
+        progress_records = [
+            p async for p in session.upload(stream=stream, size=total_size)
+        ]
+
+        assert isinstance(session.response, UploadMediaResponse)
+        assert session.response.name == "async_multi_chunk.bin"
+        assert session.response.size == total_size
+
+        phases = [p.state for p in progress_records]
+        offsets = [p.bytes_uploaded for p in progress_records]
+        assert phases == [
+            ProgressState.STARTED,
+            ProgressState.UPLOADING,
+            ProgressState.UPLOADING,
+            ProgressState.FINALIZED,
+        ]
+        assert offsets == [0, 262_144, 524_288, 600_000]
+        assert all(p.total_bytes == total_size for p in progress_records)

packages/gapic-generator/tests/system/test_resumable_upload_resume.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import io
+import os
 import pytest
 
 from google.api_core import exceptions
@@ -21,7 +22,7 @@
     ResumableUploadConfig,
     ResumableUploadSession,
 )
-from google.showcase import UploadMediaResponse
+from google.showcase import UploadMediaRequest, UploadMediaResponse
 
 from conftest import resume_resumable_upload
 
@@ -500,3 +501,110 @@ class _ClientPauseError(Exception):
     assert isinstance(session2.response, UploadMediaResponse)
     assert session2.response.name == "golden_user_resume.mp4"
     assert session2.response.size == total_size
+
+
+if os.environ.get("GAPIC_PYTHON_ASYNC", "true") == "true":
+
+    @pytest.mark.asyncio
+    async def test_async_client_upload_media_resume_direct(
+        intercepted_resumable_upload_rest_async,
+    ):
+        """Verify direct async resumption via ResumableUploadServiceAsyncClient."""
+        client, _ = intercepted_resumable_upload_rest_async
+        total_size = 1_500_000
+        chunk_size = 524_288  # 512 KiB
+        data = b"R" * total_size
+        stream = io.BytesIO(data)
+
+        class _ClientPauseError(Exception):
+            pass
+
+        session1 = await client.upload_media(
+            request=UploadMediaRequest(name="async_resume_direct.txt"),
+            config=ResumableUploadConfig(chunk_size=chunk_size),
+        )
+        with pytest.raises(_ClientPauseError):
+            async for progress in session1.upload(stream=stream, size=total_size):
+                if progress.bytes_uploaded >= chunk_size:
+                    raise _ClientPauseError("Simulated user pause after first chunk")
+
+        saved_url = session1.upload_url
+        saved_chunk_size = session1.chunk_size
+        assert saved_url is not None
+        assert saved_chunk_size == chunk_size
+
+        stream.seek(0)
+        session2 = await client.upload_media(
+            request=UploadMediaRequest(name="async_resume_direct.txt"),
+            config=ResumableUploadConfig(chunk_size=saved_chunk_size),
+        )
+        response = await session2.resume(
+            upload_url=saved_url,
+            stream=stream,
+            size=total_size,
+            chunk_size=saved_chunk_size,
+        )
+
+        assert session2.finished
+        assert isinstance(response, UploadMediaResponse)
+        assert response.name == "async_resume_direct.txt"
+        assert response.size == total_size
+
+    @pytest.mark.asyncio
+    async def test_async_client_upload_media_resume_seekable(
+        intercepted_resumable_upload_rest_async,
+    ):
+        """Verify async progress iteration during resumption via ResumableUploadServiceAsyncClient."""
+        client, _ = intercepted_resumable_upload_rest_async
+        total_size = 1_500_000
+        chunk_size = 524_288  # 512 KiB
+        data = b"G" * total_size
+        stream = io.BytesIO(data)
+
+        class _ClientPauseError(Exception):
+            pass
+
+        session1 = await client.upload_media(
+            request=UploadMediaRequest(name="async_golden_user_resume.mp4"),
+            config=ResumableUploadConfig(chunk_size=chunk_size),
+        )
+        session1_snapshots = []
+        with pytest.raises(_ClientPauseError):
+            async for progress in session1.upload(stream=stream, size=total_size):
+                session1_snapshots.append(progress)
+                if progress.bytes_uploaded >= chunk_size:
+                    raise _ClientPauseError("Simulated user pause after first chunk")
+
+        saved_url = session1.upload_url
+        saved_chunk_size = session1.chunk_size
+        assert saved_url is not None
+        assert saved_chunk_size == chunk_size
+        assert [(p.state, p.bytes_uploaded) for p in session1_snapshots] == [
+            (ProgressState.STARTED, 0),
+            (ProgressState.UPLOADING, 524_288),
+        ]
+
+        stream.seek(0)
+        session2 = await client.upload_media(
+            request=UploadMediaRequest(name="async_golden_user_resume.mp4"),
+            config=ResumableUploadConfig(chunk_size=saved_chunk_size),
+        )
+        session2_snapshots = [
+            p
+            async for p in session2.resume(
+                upload_url=saved_url,
+                stream=stream,
+                size=total_size,
+                chunk_size=saved_chunk_size,
+            )
+        ]
+
+        assert [(p.state, p.bytes_uploaded) for p in session2_snapshots] == [
+            (ProgressState.OFFSET_RECEIVED, 524_288),
+            (ProgressState.UPLOADING, 1_048_576),
+            (ProgressState.FINALIZED, 1_500_000),
+        ]
+        assert session2.finished
+        assert isinstance(session2.response, UploadMediaResponse)
+        assert session2.response.name == "async_golden_user_resume.mp4"
+        assert session2.response.size == total_size