Skip to content

Commit 954a3bd

Browse files
committed
Implement different ways on specifying key ancestry
1 parent a4b02c8 commit 954a3bd

File tree

6 files changed

+113
-14
lines changed

6 files changed

+113
-14
lines changed

appengine_fixture_loader/loader.py

+28-7
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,19 @@ def _load(od):
4343
objtype = kind[od['__kind__']]
4444
else:
4545
objtype = kind
46-
obj = objtype()
46+
47+
# set custom key if specified
48+
if '__key__' in od.keys():
49+
key = ndb.Key(*od['__key__'])
50+
obj = objtype(key=key)
51+
else:
52+
parent, id = None, None
53+
if '__parent__' in od.keys():
54+
parent = ndb.Key(*od['__parent__'])
55+
if '__id__' in od.keys():
56+
id = od['__id__']
57+
obj = objtype(parent=parent, id=id)
58+
4759
# Iterate over the non-special attributes
4860
for attribute_name in [k for k in od.keys()
4961
if not k.startswith('__') and
@@ -53,19 +65,28 @@ def _load(od):
5365
od[attribute_name])
5466
obj.__dict__['_values'][attribute_name] = attribute_value
5567

56-
# set custom key if specified
57-
if '__key__' in od.keys():
58-
obj.key = ndb.Key(*od['__key__'])
59-
6068
if post_processor:
6169
post_processor(obj)
6270

6371
# Saving obj is required to continue with the children
6472
obj.put()
6573

66-
# Scan the attributes for children
74+
# Scan the attributes for key children
75+
for child_attribute_name in [k for k in od.keys()
76+
if k == '__children__']:
77+
for o in od[child_attribute_name]:
78+
# re-create o with a new key with parent
79+
old_key = o.key
80+
new_key = ndb.Key(o.__class__, o.key.id(), parent=obj.key)
81+
o.key = new_key
82+
o.put()
83+
# delete old key
84+
old_key.delete()
85+
86+
# Scan the attributes for children properties
6787
for child_attribute_name in [k for k in od.keys()
68-
if k.startswith('__children__')]:
88+
if k.startswith('__children__')
89+
and k != '__children__']:
6990
attribute_name = child_attribute_name.split('__')[-2]
7091
for o in od[child_attribute_name]:
7192
o.__dict__['_values'][attribute_name] = obj.key

tests/ancestor_tests.py

+53-5
Original file line numberDiff line numberDiff line change
@@ -19,26 +19,74 @@ class Purchase(ndb.Model):
1919
price = ndb.IntegerProperty
2020

2121

22-
class AncestorTests(unittest.TestCase):
23-
"""Tests if we can load a JSON file"""
22+
class AncestorKeyTests(unittest.TestCase):
2423
def setUp(self):
2524
self.testbed = testbed.Testbed()
2625
self.testbed.activate()
2726
self.testbed.init_datastore_v3_stub()
2827
self.testbed.init_memcache_stub()
2928
self.customers_data = load_fixture('tests/customers.json', Customer)
30-
self.purchases_data = load_fixture('tests/purchases.json', Purchase)
29+
self.purchases_data = load_fixture('tests/purchases_key.json', Purchase)
3130

3231
def tearDown(self):
3332
self.testbed.deactivate()
3433

3534
def test_loaded_count(self):
36-
"""Make sure we got 2 objects from the JSON file"""
3735
self.assertEqual(len(self.customers_data), 2)
3836
self.assertEqual(len(self.purchases_data), 2)
3937

4038
def test_ancestors(self):
41-
"""Check whether the ancestor of keys we imported match the JSON contents"""
39+
john = Customer.query(Customer.name == 'John').get()
40+
self.assertEqual(john.name, 'John')
41+
42+
john_purchases = Purchase.query(ancestor=john.key)
43+
self.assertEqual(john_purchases.count(), 1)
44+
self.assertEqual(john_purchases.get().key.parent(), john.key)
45+
46+
47+
class AncestorParentTests(unittest.TestCase):
48+
def setUp(self):
49+
self.testbed = testbed.Testbed()
50+
self.testbed.activate()
51+
self.testbed.init_datastore_v3_stub()
52+
self.testbed.init_memcache_stub()
53+
self.customers_data = load_fixture('tests/customers.json', Customer)
54+
self.purchases_data = load_fixture('tests/purchases_parent.json', Purchase)
55+
56+
def tearDown(self):
57+
self.testbed.deactivate()
58+
59+
def test_loaded_count(self):
60+
self.assertEqual(len(self.customers_data), 2)
61+
self.assertEqual(len(self.purchases_data), 2)
62+
63+
def test_ancestors(self):
64+
john = Customer.query(Customer.name == 'John').get()
65+
self.assertEqual(john.name, 'John')
66+
67+
john_purchases = Purchase.query(ancestor=john.key)
68+
self.assertEqual(john_purchases.count(), 1)
69+
self.assertEqual(john_purchases.get().key.parent(), john.key)
70+
71+
72+
class AncestorChildrenTests(unittest.TestCase):
73+
def setUp(self):
74+
self.testbed = testbed.Testbed()
75+
self.testbed.activate()
76+
self.testbed.init_datastore_v3_stub()
77+
self.testbed.init_memcache_stub()
78+
self.data = load_fixture('tests/customers_purchases.json',
79+
{'Customer': Customer, 'Purchase': Purchase})
80+
81+
def tearDown(self):
82+
self.testbed.deactivate()
83+
84+
def test_loaded_count(self):
85+
self.assertEqual(len(self.data), 2)
86+
87+
def test_ancestors(self):
88+
self.assertEqual(Purchase.query().count(), 2)
89+
4290
john = Customer.query(Customer.name == 'John').get()
4391
self.assertEqual(john.name, 'John')
4492

tests/customers.json

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
[
22
{
33
"name": "John",
4-
"__key__": ["Customer", "john"]
4+
"__id__": "john"
55
},
66
{
77
"name": "Jane",
8-
"__key__": ["Customer", "jane"]
8+
"__id__": "jane"
99
}
1010
]

tests/customers_purchases.json

+20
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
[
2+
{
3+
"__kind__": "Customer",
4+
"name": "John",
5+
"__id__": "john",
6+
"__children__": [{
7+
"__kind__": "Purchase",
8+
"price": 100
9+
}]
10+
},
11+
{
12+
"__kind__": "Customer",
13+
"name": "Jane",
14+
"__id__": "jane",
15+
"__children__": [{
16+
"__kind__": "Purchase",
17+
"price": 50
18+
}]
19+
}
20+
]
File renamed without changes.

tests/purchases_parent.json

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
[
2+
{
3+
"price": 100,
4+
"__parent__": ["Customer", "john"]
5+
},
6+
{
7+
"price": 50,
8+
"__parent__": ["Customer", "jane"]
9+
}
10+
]

0 commit comments

Comments
 (0)