@@ -296,6 +296,89 @@ def test__handle_row_contents(self):
296296 }
297297 self .assertEqual (partial_row_data .cells , expected_cells )
298298
299+ def test_update_from_read_rows (self ):
300+ from gcloud .bigtable ._generated import bigtable_data_pb2 as data_pb2
301+ from gcloud .bigtable ._generated import (
302+ bigtable_service_messages_pb2 as messages_pb2 )
303+
304+ row_key = b'row-key'
305+ partial_row_data = self ._makeOne (row_key )
306+
307+ # Set-up chunk1, some data that will be reset by chunk2.
308+ ignored_family_name = u'ignore-name'
309+ row_contents = data_pb2 .Family (name = ignored_family_name )
310+ chunk1 = messages_pb2 .ReadRowsResponse .Chunk (row_contents = row_contents )
311+
312+ # Set-up chunk2, a reset row.
313+ chunk2 = messages_pb2 .ReadRowsResponse .Chunk (reset_row = True )
314+
315+ # Set-up chunk3, a column family with no columns.
316+ family_name = u'name'
317+ row_contents = data_pb2 .Family (name = family_name )
318+ chunk3 = messages_pb2 .ReadRowsResponse .Chunk (row_contents = row_contents )
319+
320+ # Set-up chunk4, a commit row.
321+ chunk4 = messages_pb2 .ReadRowsResponse .Chunk (commit_row = True )
322+
323+ # Prepare request and make sure PartialRowData is empty before.
324+ read_rows_response_pb = messages_pb2 .ReadRowsResponse (
325+ row_key = row_key , chunks = [chunk1 , chunk2 , chunk3 , chunk4 ])
326+ self .assertEqual (partial_row_data .cells , {})
327+ self .assertFalse (partial_row_data .committed )
328+ self .assertFalse (partial_row_data ._chunks_encountered )
329+
330+ # Parse the response and make sure the cells took place.
331+ partial_row_data .update_from_read_rows (read_rows_response_pb )
332+ self .assertEqual (partial_row_data .cells , {family_name : {}})
333+ self .assertFalse (ignored_family_name in partial_row_data .cells )
334+ self .assertTrue (partial_row_data .committed )
335+ self .assertTrue (partial_row_data ._chunks_encountered )
336+
337+ def test_update_from_read_rows_while_committed (self ):
338+ partial_row_data = self ._makeOne (None )
339+ partial_row_data ._committed = True
340+ self .assertFalse (partial_row_data ._chunks_encountered )
341+
342+ with self .assertRaises (ValueError ):
343+ partial_row_data .update_from_read_rows (None )
344+
345+ self .assertFalse (partial_row_data ._chunks_encountered )
346+
347+ def test_update_from_read_rows_row_key_disagree (self ):
348+ from gcloud .bigtable ._generated import (
349+ bigtable_service_messages_pb2 as messages_pb2 )
350+
351+ row_key1 = b'row-key1'
352+ row_key2 = b'row-key2'
353+ partial_row_data = self ._makeOne (row_key1 )
354+ self .assertFalse (partial_row_data ._chunks_encountered )
355+
356+ self .assertNotEqual (row_key1 , row_key2 )
357+ read_rows_response_pb = messages_pb2 .ReadRowsResponse (row_key = row_key2 )
358+ with self .assertRaises (ValueError ):
359+ partial_row_data .update_from_read_rows (read_rows_response_pb )
360+
361+ self .assertFalse (partial_row_data ._chunks_encountered )
362+
363+ def test_update_from_read_rows_empty_chunk (self ):
364+ from gcloud .bigtable ._generated import (
365+ bigtable_service_messages_pb2 as messages_pb2 )
366+
367+ row_key = b'row-key'
368+ partial_row_data = self ._makeOne (row_key )
369+ self .assertFalse (partial_row_data ._chunks_encountered )
370+
371+ chunk = messages_pb2 .ReadRowsResponse .Chunk ()
372+ read_rows_response_pb = messages_pb2 .ReadRowsResponse (
373+ row_key = row_key , chunks = [chunk ])
374+
375+ # This makes it an "empty" chunk.
376+ self .assertEqual (chunk .WhichOneof ('chunk' ), None )
377+ with self .assertRaises (ValueError ):
378+ partial_row_data .update_from_read_rows (read_rows_response_pb )
379+
380+ self .assertFalse (partial_row_data ._chunks_encountered )
381+
299382
300383class TestPartialRowsData (unittest2 .TestCase ):
301384
0 commit comments