@@ -38,7 +38,9 @@ def test_ctor(self):
38
38
self .assertEqual (batch .namespace , _NAMESPACE )
39
39
self .assertTrue (batch ._id is None )
40
40
self .assertEqual (batch ._status , batch ._INITIAL )
41
- self .assertTrue (isinstance (batch .mutations , datastore_pb2 .Mutation ))
41
+ self .assertTrue (isinstance (batch ._commit_request ,
42
+ datastore_pb2 .CommitRequest ))
43
+ self .assertTrue (batch .mutations is batch ._commit_request .mutations )
42
44
self .assertEqual (batch ._partial_key_entities , [])
43
45
44
46
def test_current (self ):
@@ -90,7 +92,7 @@ def test_put_entity_w_partial_key(self):
90
92
91
93
batch .put (entity )
92
94
93
- mutated_entity = _mutated_pb (self , batch .mutations , 'insert_auto_id ' )
95
+ mutated_entity = _mutated_pb (self , batch .mutations , 'insert ' )
94
96
self .assertEqual (mutated_entity .key , key ._key )
95
97
self .assertEqual (batch ._partial_key_entities , [entity ])
96
98
@@ -121,9 +123,10 @@ def test_put_entity_w_completed_key(self):
121
123
self .assertFalse (prop_dict ['foo' ].exclude_from_indexes )
122
124
self .assertTrue (prop_dict ['baz' ].exclude_from_indexes )
123
125
self .assertFalse (prop_dict ['spam' ].exclude_from_indexes )
124
- self .assertTrue (prop_dict ['spam' ].list_value [0 ].exclude_from_indexes )
125
- self .assertTrue (prop_dict ['spam' ].list_value [1 ].exclude_from_indexes )
126
- self .assertTrue (prop_dict ['spam' ].list_value [2 ].exclude_from_indexes )
126
+ spam_values = prop_dict ['spam' ].array_value .values
127
+ self .assertTrue (spam_values [0 ].exclude_from_indexes )
128
+ self .assertTrue (spam_values [1 ].exclude_from_indexes )
129
+ self .assertTrue (spam_values [2 ].exclude_from_indexes )
127
130
self .assertFalse ('frotz' in prop_dict )
128
131
129
132
def test_put_entity_w_completed_key_prefixed_project (self ):
@@ -153,9 +156,10 @@ def test_put_entity_w_completed_key_prefixed_project(self):
153
156
self .assertFalse (prop_dict ['foo' ].exclude_from_indexes )
154
157
self .assertTrue (prop_dict ['baz' ].exclude_from_indexes )
155
158
self .assertFalse (prop_dict ['spam' ].exclude_from_indexes )
156
- self .assertTrue (prop_dict ['spam' ].list_value [0 ].exclude_from_indexes )
157
- self .assertTrue (prop_dict ['spam' ].list_value [1 ].exclude_from_indexes )
158
- self .assertTrue (prop_dict ['spam' ].list_value [2 ].exclude_from_indexes )
159
+ spam_values = prop_dict ['spam' ].array_value .values
160
+ self .assertTrue (spam_values [0 ].exclude_from_indexes )
161
+ self .assertTrue (spam_values [1 ].exclude_from_indexes )
162
+ self .assertTrue (spam_values [2 ].exclude_from_indexes )
159
163
self .assertFalse ('frotz' in prop_dict )
160
164
161
165
def test_delete_w_partial_key (self ):
@@ -424,20 +428,18 @@ def current_batch(self):
424
428
return self ._batches [0 ]
425
429
426
430
427
- def _assert_num_mutations (test_case , mutation_pb , num_mutations ):
428
- total_mutations = (len (mutation_pb .upsert ) +
429
- len (mutation_pb .update ) +
430
- len (mutation_pb .insert ) +
431
- len (mutation_pb .insert_auto_id ) +
432
- len (mutation_pb .delete ))
433
- test_case .assertEqual (total_mutations , num_mutations )
431
+ def _assert_num_mutations (test_case , mutation_pb_list , num_mutations ):
432
+ test_case .assertEqual (len (mutation_pb_list ), num_mutations )
434
433
435
434
436
- def _mutated_pb (test_case , mutation_pb , mutation_type ):
435
+ def _mutated_pb (test_case , mutation_pb_list , mutation_type ):
437
436
# Make sure there is only one mutation.
438
- _assert_num_mutations (test_case , mutation_pb , 1 )
437
+ _assert_num_mutations (test_case , mutation_pb_list , 1 )
439
438
440
- mutated_pbs = getattr (mutation_pb , mutation_type , [])
441
- # Make sure we have exactly one protobuf.
442
- test_case .assertEqual (len (mutated_pbs ), 1 )
443
- return mutated_pbs [0 ]
439
+ # We grab the only mutation.
440
+ mutated_pb = mutation_pb_list [0 ]
441
+ # Then check if it is the correct type.
442
+ test_case .assertEqual (mutated_pb .WhichOneof ('operation' ),
443
+ mutation_type )
444
+
445
+ return getattr (mutated_pb , mutation_type )
0 commit comments