Skip to content

Commit e56a275

Browse files
KamiKami
committed
Add tests which exercise new API methods which work with raw Entity
Protobuf objects.
1 parent c334a2b commit e56a275

1 file changed

Lines changed: 147 additions & 0 deletions

File tree

‎datastore/tests/system/test_system.py‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from google.cloud.datastore.helpers import GeoPoint
2525
from google.cloud.environment_vars import GCD_DATASET
2626
from google.cloud.exceptions import Conflict
27+
from google.cloud.datastore_v1.proto import entity_pb2
2728

2829
from test_utils.system import EmulatorCreds
2930
from test_utils.system import unique_resource_id
@@ -554,3 +555,149 @@ def test_empty_array_put(self):
554555
retrieved = local_client.get(entity.key)
555556

556557
self.assertEqual(entity["children"], retrieved["children"])
558+
559+
560+
class TestDatastoreRawEntityPBMethods(TestDatastore):
561+
"""
562+
Test which exercise methods which works with the raw Entity Protobof objects.
563+
"""
564+
565+
def test_put_and_get_entity_pb(self):
566+
# 1. Store the Entity
567+
key = Config.CLIENT.key("TestEntityPB", 1111)
568+
entity_pb = entity_pb2.Entity()
569+
entity_pb.key.CopyFrom(key.to_protobuf())
570+
571+
int32_value = entity_pb.properties.get_or_create('int32_key')
572+
int32_value.integer_value = 1000
573+
string_value = entity_pb.properties.get_or_create('string_key')
574+
string_value.string_value = 'foo bar baz'
575+
576+
Config.CLIENT.put_entity_pb(entity_pb)
577+
578+
# 2. Verify that it was stored correctly
579+
entity_retrieved_pb = Config.CLIENT.get_entity_pb(key=key)
580+
self.assertEqual(entity_retrieved_pb, entity_pb)
581+
self.assertEqual(entity_retrieved_pb.key.path[0].kind, "TestEntityPB")
582+
self.assertEqual(entity_retrieved_pb.key.path[0].id, 1111)
583+
self.assertEqual(entity_retrieved_pb.properties['int32_key'].integer_value,
584+
1000)
585+
self.assertEqual(entity_retrieved_pb.properties['string_key'].string_value,
586+
'foo bar baz')
587+
588+
def test_put_multi_get_multi_pb(self):
589+
# Store multiple entities
590+
key1 = Config.CLIENT.key("TestEntityPB", 1111)
591+
entity1_pb = entity_pb2.Entity()
592+
entity1_pb.key.CopyFrom(key1.to_protobuf())
593+
594+
key2 = Config.CLIENT.key("TestEntityPB", 2222)
595+
entity2_pb = entity_pb2.Entity()
596+
entity2_pb.key.CopyFrom(key2.to_protobuf())
597+
598+
entity_pbs = [entity1_pb, entity2_pb]
599+
Config.CLIENT.put_multi_entity_pbs(entity_pbs=entity_pbs)
600+
601+
# Verify they were stored correctly
602+
keys = [key1, key2]
603+
entity_retrieved_pbs = Config.CLIENT.get_multi_entity_pb(keys=keys)
604+
self.assertEqual(len(entity_retrieved_pbs), 2)
605+
self.assertEqual(entity_retrieved_pbs[0].key.path[0].id, 1111)
606+
self.assertEqual(entity_retrieved_pbs[1].key.path[0].id, 2222)
607+
608+
def test_fetch_entity_pbs(self):
609+
# Store multiple entities
610+
key1 = Config.CLIENT.key("TestEntityPB", 1111)
611+
entity1_pb = entity_pb2.Entity()
612+
entity1_pb.key.CopyFrom(key1.to_protobuf())
613+
614+
int32_value = entity1_pb.properties.get_or_create('int32_key')
615+
int32_value.integer_value = 1000
616+
string_value = entity1_pb.properties.get_or_create('string_key')
617+
string_value.string_value = 'foo1'
618+
619+
key2 = Config.CLIENT.key("TestEntityPB", 2222)
620+
entity2_pb = entity_pb2.Entity()
621+
entity2_pb.key.CopyFrom(key2.to_protobuf())
622+
623+
int32_value = entity2_pb.properties.get_or_create('int32_key')
624+
int32_value.integer_value = 2000
625+
string_value = entity2_pb.properties.get_or_create('string_key')
626+
string_value.string_value = 'foo2'
627+
628+
key3 = Config.CLIENT.key("TestEntityPB", 3333)
629+
entity3_pb = entity_pb2.Entity()
630+
entity3_pb.key.CopyFrom(key3.to_protobuf())
631+
632+
int32_value = entity3_pb.properties.get_or_create('int32_key')
633+
int32_value.integer_value = 3000
634+
string_value = entity3_pb.properties.get_or_create('string_key')
635+
string_value.string_value = 'foo2'
636+
637+
key4 = Config.CLIENT.key("TestEntityPB", 44444)
638+
entity4_pb = entity_pb2.Entity()
639+
entity4_pb.key.CopyFrom(key4.to_protobuf())
640+
641+
int32_value = entity4_pb.properties.get_or_create('int32_key')
642+
int32_value.integer_value = 3000
643+
string_value = entity4_pb.properties.get_or_create('string_key')
644+
string_value.string_value = 'foo4'
645+
646+
entity_pbs = [entity1_pb, entity2_pb, entity3_pb, entity4_pb]
647+
Config.CLIENT.put_multi_entity_pbs(entity_pbs=entity_pbs)
648+
649+
# 1. Simply retrieve all
650+
query = Config.CLIENT.query(kind='TestEntityPB')
651+
entity_retrieved_pbs = list(query.fetch_entity_pb())
652+
653+
self.assertEqual(len(entity_retrieved_pbs), 4)
654+
self.assertEqual(entity_retrieved_pbs, entity_pbs)
655+
656+
# 2. Simple query on integer field
657+
query = Config.CLIENT.query(kind='TestEntityPB')
658+
query.add_filter('int32_key', '>=', 2000)
659+
entity_retrieved_pbs = list(query.fetch_entity_pb())
660+
661+
self.assertEqual(len(entity_retrieved_pbs), 3)
662+
self.assertEqual(entity_retrieved_pbs[0].properties['int32_key']
663+
.integer_value,
664+
2000)
665+
self.assertEqual(entity_retrieved_pbs[1].properties['int32_key']
666+
.integer_value,
667+
3000)
668+
self.assertEqual(entity_retrieved_pbs[2].properties['int32_key']
669+
.integer_value,
670+
3000)
671+
672+
query = Config.CLIENT.query(kind='TestEntityPB')
673+
query.add_filter('int32_key', '>', 3000)
674+
entity_retrieved_pbs = list(query.fetch_entity_pb())
675+
676+
self.assertEqual(len(entity_retrieved_pbs), 0)
677+
678+
# 3. Simple query on string field
679+
query = Config.CLIENT.query(kind='TestEntityPB')
680+
query.add_filter('string_key', '=', 'foo2')
681+
entity_retrieved_pbs = list(query.fetch_entity_pb())
682+
683+
self.assertEqual(len(entity_retrieved_pbs), 2)
684+
self.assertEqual(entity_retrieved_pbs[0].properties['string_key']
685+
.string_value,
686+
'foo2')
687+
self.assertEqual(entity_retrieved_pbs[1].properties['string_key']
688+
.string_value,
689+
'foo2')
690+
691+
# 4. Complex query on integer and string field value
692+
query = Config.CLIENT.query(kind='TestEntityPB')
693+
query.add_filter('int32_key', '=', 3000)
694+
query.add_filter('string_key', '=', 'foo2')
695+
entity_retrieved_pbs = list(query.fetch_entity_pb())
696+
697+
self.assertEqual(len(entity_retrieved_pbs), 1)
698+
self.assertEqual(entity_retrieved_pbs[0].properties['int32_key']
699+
.integer_value,
700+
3000)
701+
self.assertEqual(entity_retrieved_pbs[0].properties['string_key']
702+
.string_value,
703+
'foo2')

0 commit comments

Comments
 (0)