resolve comments

Author: Pulkit0110

google/cloud/storage/_experimental/asyncio/retry/writes_resumption_strategy.py

@@ -37,7 +37,7 @@ class _WriteState:
 
     def __init__(
         self,
-        spec: storage_type.AppendObjectSpec,
+        spec: Union[storage_type.AppendObjectSpec, storage_type.WriteObjectSpec],
         chunk_size: int,
         user_buffer: IO[bytes],
     ):
@@ -48,7 +48,7 @@ def __init__(
         self.bytes_sent: int = 0
         self.write_handle: Union[bytes, storage_type.BidiWriteHandle, None] = None
         self.routing_token: Optional[str] = None
-        self.is_complete: bool = False
+        self.is_finalized: bool = False
 
 
 class _WriteResumptionStrategy(_BaseResumptionStrategy):
@@ -61,34 +61,29 @@ def generate_requests(
 
         For Appendable Objects, every stream opening should send an
         AppendObjectSpec. If resuming, the `write_handle` is added to that spec.
+
+        This method is not applicable for `open` methods.
         """
         write_state: _WriteState = state["write_state"]
 
-        if write_state.write_handle:
-            write_state.spec.write_handle = write_state.write_handle
-
-        if write_state.routing_token:
-            write_state.spec.routing_token = write_state.routing_token
-
-        # Initial request of the stream must provide the specification.
-        # If we have a write_handle, we request a state lookup to verify persisted offset.
-        do_state_lookup = write_state.write_handle is not None
+        initial_request = storage_type.BidiWriteObjectRequest()
 
         # Determine if we need to send WriteObjectSpec or AppendObjectSpec
-        initial_request = storage_type.BidiWriteObjectRequest(
-            state_lookup=do_state_lookup
-        )
-
         if isinstance(write_state.spec, storage_type.WriteObjectSpec):
             initial_request.write_object_spec = write_state.spec
         else:
+            if write_state.write_handle:
+                write_state.spec.write_handle = write_state.write_handle
+
+            if write_state.routing_token:
+                write_state.spec.routing_token = write_state.routing_token
             initial_request.append_object_spec = write_state.spec
 
         yield initial_request
 
         # The buffer should already be seeked to the correct position (persisted_size)
         # by the `recover_state_on_failure` method before this is called.
-        while not write_state.is_complete:
+        while not write_state.is_finalized:
             chunk = write_state.user_buffer.read(write_state.chunk_size)
 
             # End of File detection
@@ -113,16 +108,16 @@ def update_state_from_response(
         """Processes a server response and updates the write state."""
         write_state: _WriteState = state["write_state"]
 
-        if response.persisted_size is not None:
-            if response.persisted_size > write_state.persisted_size:
-                write_state.persisted_size = response.persisted_size
+        if response.persisted_size:
+            write_state.persisted_size = response.persisted_size
 
         if response.write_handle:
             write_state.write_handle = response.write_handle
 
         if response.resource:
-            write_state.is_complete = True
             write_state.persisted_size = response.resource.size
+            if response.resource.finalize_time:
+                write_state.is_finalized = True
 
     async def recover_state_on_failure(
         self, error: Exception, state: Dict[str, Any]

tests/unit/asyncio/retry/test_writes_resumption_strategy.py

@@ -15,6 +15,7 @@
 import io
 import unittest
 import unittest.mock as mock
+from datetime import datetime
 
 import pytest
 import google_crc32c
@@ -137,7 +138,6 @@ def test_generate_requests_resumption(self):
         self.assertEqual(
             requests[0].append_object_spec.write_handle.handle, b"test-handle"
         )
-        self.assertTrue(requests[0].state_lookup)
 
         # Check data starts from offset 4
         self.assertEqual(requests[1].write_offset, 4)
@@ -206,28 +206,11 @@ def test_update_state_from_response(self):
         strategy.update_state_from_response(response2, state)
         self.assertEqual(write_state.persisted_size, 1024)
 
-        final_resource = storage_type.Object(name="test-object", bucket="b", size=2048)
+        final_resource = storage_type.Object(name="test-object", bucket="b", size=2048, finalize_time=datetime.now())
         response3 = storage_type.BidiWriteObjectResponse(resource=final_resource)
         strategy.update_state_from_response(response3, state)
-        self.assertTrue(write_state.is_complete)
-        self.assertEqual(write_state.persisted_size, 2048)
-
-    def test_update_state_from_response_ignores_smaller_persisted_size(self):
-        strategy = self._make_one()
-        state = {
-            "write_state": _WriteState(
-                mock.Mock(spec=storage_type.AppendObjectSpec),
-                0,
-                mock.Mock(spec=io.BytesIO),
-            ),
-        }
-        write_state = state["write_state"]
-        write_state.persisted_size = 2048
-
-        response = storage_type.BidiWriteObjectResponse(persisted_size=1024)
-        strategy.update_state_from_response(response, state)
-
         self.assertEqual(write_state.persisted_size, 2048)
+        self.assertTrue(write_state.is_finalized)
 
     @pytest.mark.asyncio
     async def test_recover_state_on_failure_handles_redirect(self):