fix unit test

Author: Linchin

samples/pyarrow/append_rows_with_arrow.py

@@ -159,16 +159,14 @@ def generate_write_request_with_pyarrow(num_rows=TABLE_LENGTH):
     # Determine max_chunksize of the record batches. Because max size of
     # AppendRowsRequest is 10 MB, we need to split the table if it's too big.
     # See: https://github.com/googleapis/googleapis/blob/27296636cf8797026124cd67034b42190ab602a4/google/cloud/bigquery/storage/v1/storage.proto#L422
-    max_request_bytes = 10 * 2**20   # 10 MB
+    max_request_bytes = 10 * 2**20  # 10 MB
     chunk_num = int(table.nbytes / max_request_bytes) + 1
     chunk_size = int(table.num_rows / chunk_num)
 
     # Construct request(s).
     for batch in table.to_batches(max_chunksize=chunk_size):
         request = gapic_types.AppendRowsRequest()
-        request.arrow_rows.rows.serialized_record_batch = (
-            batch.serialize().to_pybytes()
-        )
+        request.arrow_rows.rows.serialized_record_batch = batch.serialize().to_pybytes()
         yield request
 
 

tests/unit/test_writer_v1.py

@@ -203,12 +203,12 @@ def _make_mock_client():
         return mock.create_autospec(big_query_write.BigQueryWriteClient)
 
     @staticmethod
-    def _make_mock_stream():
+    def _make_mock_stream(initial_template=REQUEST_TEMPLATE):
         from google.cloud.bigquery_storage_v1.writer import _process_request_template
 
         writer = mock.Mock()
         template = mock.PropertyMock(
-            return_value=_process_request_template(REQUEST_TEMPLATE)
+            return_value=_process_request_template(initial_template)
         )
         type(writer)._initial_request_template = template
         return writer
@@ -251,8 +251,17 @@ def test_is_active(self):
     def test_initial_send(self, background_consumer, bidi_rpc):
         from google.cloud.bigquery_storage_v1.writer import AppendRowsFuture
 
+        initial_request_template = gapic_types.AppendRowsRequest(
+            write_stream="stream-name-from-REQUEST_TEMPLATE",
+            offset=0,
+            proto_rows=gapic_types.AppendRowsRequest.ProtoData(
+                writer_schema=gapic_types.ProtoSchema(
+                    proto_descriptor=descriptor_pb2.DescriptorProto()
+                )
+            ),
+        )
         mock_client = self._make_mock_client()
-        mock_stream = self._make_mock_stream()
+        mock_stream = self._make_mock_stream(initial_template=initial_request_template)
         connection = self._make_one(mock_client, mock_stream)
 
         type(bidi_rpc.return_value).is_active = mock.PropertyMock(