#!/usr/bin/env vpython # Copyright 2014 The LUCI Authors. All rights reserved. # Use of this source code is governed under the Apache License, Version 2.0 # that can be found in the LICENSE file. import sys import unittest from test_support import test_env test_env.setup_test_env() from google.appengine.ext import ndb from components.datastore_utils import monotonic from components.datastore_utils import txn from test_support import test_case # Access to a protected member _XX of a client class - pylint: disable=W0212 class EntityX(ndb.Model): a = ndb.IntegerProperty() def _pre_put_hook(self): super(EntityX, self)._pre_put_hook() class EntityY(ndb.Model): def _pre_put_hook(self): super(EntityY, self)._pre_put_hook() class MonotonicTest(test_case.TestCase): def setUp(self): super(MonotonicTest, self).setUp() self.parent = ndb.Key('Root', 1) def test_insert(self): data = EntityX(id=1, parent=self.parent) called = [] self.mock(EntityX, '_pre_put_hook', lambda _: called.append(1)) actual = monotonic.insert(data, None) expected = ndb.Key('EntityX', 1, parent=self.parent) self.assertEqual(expected, actual) self.assertEqual([1], called) def test_insert_already_present(self): EntityX(id=1, parent=self.parent).put() data = EntityX(id=1, parent=self.parent) actual = monotonic.insert(data, None) self.assertEqual(None, actual) def test_insert_new_key(self): data = EntityX(id=1, parent=self.parent) extra = EntityY(id=1, parent=data.key) # Make sure the _pre_put_hook functions are called. called = [] self.mock(EntityX, '_pre_put_hook', lambda _: called.append(1)) self.mock(EntityY, '_pre_put_hook', lambda _: called.append(2)) actual = monotonic.insert(data, self.fail, extra=[extra]) expected = ndb.Key('EntityX', 1, parent=self.parent) self.assertEqual(expected, actual) self.assertEqual([1, 2], called) def test_insert_new_key_already_present(self): EntityX(id=1, parent=self.parent).put() data = EntityX(id=1, parent=self.parent) called = [] self.mock(EntityX, '_pre_put_hook', lambda _: called.append(1)) new_key = ndb.Key('EntityX', 2, parent=self.parent) actual = monotonic.insert(data, lambda: called.append(2) or new_key) expected = ndb.Key('EntityX', 2, parent=self.parent) self.assertEqual(expected, actual) self.assertEqual([2, 1], called) def test_insert_new_key_already_present_twice(self): EntityX(id=1, parent=self.parent).put() EntityX(id=2, parent=self.parent).put() data = EntityX(id=1, parent=self.parent) new_keys = [ ndb.Key('EntityX', 2, parent=self.parent), ndb.Key('EntityX', 3, parent=self.parent), ] actual = monotonic.insert(data, lambda: new_keys.pop(0)) self.assertEqual([], new_keys) expected = ndb.Key('EntityX', 3, parent=self.parent) self.assertEqual(expected, actual) def test_insert_new_key_already_present_twice_fail_after(self): EntityX(id=1, parent=self.parent).put() EntityX(id=2, parent=self.parent).put() EntityX(id=3, parent=self.parent).put() data = EntityX(id=1, parent=self.parent) new_keys = [ ndb.Key('EntityX', 2, parent=self.parent), ndb.Key('EntityX', 3, parent=self.parent), ] actual = monotonic.insert( data, lambda: new_keys.pop(0) if new_keys else None) self.assertEqual([], new_keys) self.assertEqual(None, actual) def test_insert_transaction_failure(self): EntityX(id=1, parent=self.parent).put() calls = [] def transaction_async(*args, **kwargs): calls.append(1) if len(calls) < 2: raise txn.CommitError() return old_transaction_async(*args, **kwargs) old_transaction_async = self.mock( txn, 'transaction_async', transaction_async) actual = monotonic.insert(EntityX(id=2, parent=self.parent)) expected = ndb.Key('EntityX', 2, parent=self.parent) self.assertEqual(expected, actual) self.assertEqual([1, 1], calls) def test_get_versioned_root_model(self): cls = monotonic.get_versioned_root_model('fidoula') self.assertEqual('fidoula', cls._get_kind()) self.assertTrue(issubclass(cls, ndb.Model)) self.assertEqual(53, cls(current=53).current) def test_get_versioned_most_recent(self): # First entity id is HIGH_KEY_ID, second is HIGH_KEY_ID-1. cls = monotonic.get_versioned_root_model('fidoula') parent_key = ndb.Key(cls, 'foo') for i in (monotonic.HIGH_KEY_ID, monotonic.HIGH_KEY_ID-1): monotonic.store_new_version(EntityX(parent=parent_key), cls) actual = monotonic.get_versioned_most_recent(EntityX, parent_key) expected = EntityX(key=ndb.Key('EntityX', i, parent=parent_key)) self.assertEqual(expected, actual) def test_get_versioned_most_recent_with_root(self): # First entity id is HIGH_KEY_ID, second is HIGH_KEY_ID-1. cls = monotonic.get_versioned_root_model('fidoula') parent_key = ndb.Key(cls, 'foo') for i in (monotonic.HIGH_KEY_ID, monotonic.HIGH_KEY_ID-1): monotonic.store_new_version(EntityX(parent=parent_key), cls) actual = monotonic.get_versioned_most_recent_with_root( EntityX, parent_key) expected = ( cls(key=parent_key, current=i), EntityX(key=ndb.Key('EntityX', i, parent=parent_key)), ) self.assertEqual(expected, actual) def test_get_versioned_most_recent_with_root_already_saved(self): # Stores the root entity with .current == None. cls = monotonic.get_versioned_root_model('fidoula') parent_key = ndb.Key(cls, 'foo') cls(key=parent_key).put() monotonic.store_new_version(EntityX(parent=parent_key), cls) actual = monotonic.get_versioned_most_recent_with_root(EntityX, parent_key) expected = ( cls(key=parent_key, current=monotonic.HIGH_KEY_ID), EntityX(key=ndb.Key('EntityX', monotonic.HIGH_KEY_ID, parent=parent_key)), ) self.assertEqual(expected, actual) def test_get_versioned_most_recent_with_root_already_saved_invalid(self): # Stores the root entity with an invalid .current value. cls = monotonic.get_versioned_root_model('fidoula') parent_key = ndb.Key(cls, 'foo') cls(key=parent_key, current=23).put() monotonic.store_new_version(EntityX(parent=parent_key), cls) actual = monotonic.get_versioned_most_recent_with_root(EntityX, parent_key) expected = ( cls(key=parent_key, current=23), EntityX(key=ndb.Key('EntityX', 23, parent=parent_key)), ) self.assertEqual(expected, actual) def test_get_versioned_most_recent_with_root_unexpected_extra(self): cls = monotonic.get_versioned_root_model('fidoula') parent_key = ndb.Key(cls, 'foo') monotonic.store_new_version(EntityX(parent=parent_key), cls) monotonic.store_new_version(EntityX(parent=parent_key), cls) EntityX(id=monotonic.HIGH_KEY_ID-2, parent=parent_key).put() # The unexpected entity is not registered. actual = monotonic.get_versioned_most_recent_with_root(EntityX, parent_key) expected = ( cls(key=parent_key, current=monotonic.HIGH_KEY_ID-1), EntityX( key=ndb.Key('EntityX', monotonic.HIGH_KEY_ID-1, parent=parent_key)), ) self.assertEqual(expected, actual) # The unexpected entity is safely skipped. In particular, root.current was # updated properly. monotonic.store_new_version(EntityX(parent=parent_key), cls) actual = monotonic.get_versioned_most_recent_with_root(EntityX, parent_key) expected = ( cls(key=parent_key, current=monotonic.HIGH_KEY_ID-3), EntityX( key=ndb.Key('EntityX', monotonic.HIGH_KEY_ID-3, parent=parent_key)), ) self.assertEqual(expected, actual) def test_store_new_version(self): cls = monotonic.get_versioned_root_model('fidoula') parent = ndb.Key(cls, 'foo') actual = monotonic.store_new_version(EntityX(a=1, parent=parent), cls) self.assertEqual( ndb.Key('fidoula', 'foo', 'EntityX', monotonic.HIGH_KEY_ID), actual) actual = monotonic.store_new_version(EntityX(a=2, parent=parent), cls) self.assertEqual( ndb.Key('fidoula', 'foo', 'EntityX', monotonic.HIGH_KEY_ID - 1), actual) def test_store_new_version_extra(self): # Includes an unrelated entity in the PUT. It must be in the same entity # group. cls = monotonic.get_versioned_root_model('fidoula') parent = ndb.Key(cls, 'foo') class Unrelated(ndb.Model): b = ndb.IntegerProperty() unrelated = Unrelated(id='bar', parent=parent, b=42) actual = monotonic.store_new_version( EntityX(a=1, parent=parent), cls, extra=[unrelated]) self.assertEqual( ndb.Key('fidoula', 'foo', 'EntityX', monotonic.HIGH_KEY_ID), actual) actual = monotonic.store_new_version(EntityX(a=2, parent=parent), cls) self.assertEqual( ndb.Key('fidoula', 'foo', 'EntityX', monotonic.HIGH_KEY_ID - 1), actual) self.assertEqual({'b': 42}, unrelated.key.get().to_dict()) def test_store_new_version_transaction_failure(self): # Ensures that when a transaction fails, the key id is not modified and the # retry is on the same key id. cls = monotonic.get_versioned_root_model('fidoula') parent = ndb.Key(cls, 'foo') actual = monotonic.store_new_version(EntityX(a=1, parent=parent), cls) calls = [] def transaction_async(*args, **kwargs): calls.append(1) if len(calls) < 2: raise txn.CommitError() return old_transaction_async(*args, **kwargs) old_transaction_async = self.mock( txn, 'transaction_async', transaction_async) actual = monotonic.store_new_version(EntityX(a=2, parent=parent), cls) self.assertEqual( ndb.Key('fidoula', 'foo', 'EntityX', monotonic.HIGH_KEY_ID - 1), actual) self.assertEqual([1, 1], calls) if __name__ == '__main__': if '-v' in sys.argv: unittest.TestCase.maxDiff = None unittest.main()