[Findit] Extending versioned_model.py to support versioning multiple entities of the same class.

appengine/findit/model/test/versioned_model_test.py

 # Copyright 2015 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 from google.appengine.api import datastore_errors
 from google.appengine.ext import ndb
 
 from testing_utils import testing
 
 from model.versioned_model import VersionedModel
 
 
-class _Entity(VersionedModel):
+class _SingularEntity(VersionedModel):
+  value = ndb.IntegerProperty(indexed=False)
+
+
+class _MultipleEntity(VersionedModel):
+
+  @staticmethod
+  def Create(string_id):
+    return _MultipleEntity(key=ndb.Key('_MultipleEntity', string_id))
+
   value = ndb.IntegerProperty(indexed=False)
 
 
 class VersionedModelTest(testing.AppengineTestCase):
 
   def testGetRootModel(self):
-    root_model_class = _Entity._GetRootModel()
-    self.assertEqual('_EntityRoot', root_model_class._get_kind())
+    root_model_class = _SingularEntity._GetRootModel()
+    self.assertEqual('_SingularEntityRoot', root_model_class._get_kind())
     self.assertTrue(issubclass(root_model_class, ndb.Model))
     self.assertEqual(3, root_model_class(current=3).current)
 
   def testDefaultVersionIsZero(self):
-    entity = _Entity()
-    self.assertEqual(0, entity.version)
+    self.assertEqual(0, _SingularEntity().version)
+    self.assertEqual(0, _MultipleEntity().version)
 
-  def testGetMostRecentVersionWhenNoData(self):
-    entity = _Entity.GetVersion()
-    self.assertIsNone(entity)
+  def testVersionForMultipleEntity(self):
+    multiple_entity = _MultipleEntity.Create('m1')
+    self.assertEqual(0, multiple_entity.version)
 
-  def testGetMostRecentVersionWhenDataExists(self):
-    root_key = ndb.Key('_EntityRoot', 1)
-    _Entity._GetRootModel()(key=root_key, current=2).put()
-    _Entity(key=ndb.Key('_Entity', 1, parent=root_key), value=1).put()
-    _Entity(key=ndb.Key('_Entity', 2, parent=root_key), value=2).put()
+  def testGetMostRecentVersionWhenDataExistsForSingularEntity(self):
+    root_key = ndb.Key('_SingularEntityRoot', 1)
+    _SingularEntity._GetRootModel()(key=root_key, current=2).put()
+    _SingularEntity(
+        key=ndb.Key('_SingularEntity', 1, parent=root_key), value=1).put()
+    _SingularEntity(
+        key=ndb.Key('_SingularEntity', 2, parent=root_key), value=2).put()
 
-    entity = _Entity.GetVersion()
+    entity = _SingularEntity.GetVersion()
     self.assertEqual(2, entity.version)
     self.assertEqual(2, entity.value)
 
-  def testGetNextVersionWhenDataExists(self):
-    root_key = ndb.Key('_EntityRoot', 1)
-    _Entity._GetRootModel()(key=root_key, current=2).put()
-    _Entity(key=ndb.Key('_Entity', 1, parent=root_key), value=1).put()
-    _Entity(key=ndb.Key('_Entity', 2, parent=root_key), value=2).put()
+  def testGetMostRecentVersionWhenNoData(self):
+    self.assertIsNone(_SingularEntity.GetVersion())
+    self.assertIsNone(_MultipleEntity.Create('m1').GetVersion())
 
-    entity = _Entity.GetVersion(2)
+  def testCreateMultipleEntity(self):
+    # Tests creating multiple versioned entities of the same model type.
+    multiple_entity_1 = _MultipleEntity.Create('m1')
+    multiple_entity_1.value = 1
+    multiple_entity_1.Save()
+
+    multiple_entity_2 = _MultipleEntity.Create('m2')
+    multiple_entity_2.value = 2
+    multiple_entity_2.Save()
+
+    self.assertEqual(multiple_entity_1.version, 1)
+    self.assertEqual(multiple_entity_1.value, 1)
+    self.assertEqual(multiple_entity_2.version, 1)
+    self.assertEqual(multiple_entity_2.value, 2)
+
+  def testUpdateMultipleEntity(self):
+    multiple_entity = _MultipleEntity.Create('m1')
+    multiple_entity.value = 1
+    multiple_entity.Save()
+    multiple_entity.value = 3
+    multiple_entity.Save()
+
+    self.assertEqual(multiple_entity.version, 2)
+    self.assertEqual(multiple_entity.value, 3)
+
+  def testGetNextVersionWhenDataExistsForSingularEntity(self):
+    root_key = ndb.Key('_SingularEntityRoot', 1)
+    _SingularEntity._GetRootModel()(key=root_key, current=2).put()
+    _SingularEntity(
+        key=ndb.Key('_SingularEntity', 1, parent=root_key), value=1).put()
+    _SingularEntity(
+        key=ndb.Key('_SingularEntity', 2, parent=root_key), value=2).put()
+
+    entity = _SingularEntity.GetVersion(version=2)
     self.assertEqual(2, entity.version)
     self.assertEqual(2, entity.value)
-    self.assertIsNone(_Entity.GetVersion(0))
-    self.assertIsNone(_Entity.GetVersion(3))
+    self.assertIsNone(_SingularEntity.GetVersion(version=0))
+    self.assertIsNone(_SingularEntity.GetVersion(version=3))
 
-  def testGetLatestVersionNumber(self):
-    root_key = ndb.Key('_EntityRoot', 1)
-    _Entity._GetRootModel()(key=root_key, current=1).put()
-    _Entity(key=ndb.Key('_Entity', 1, parent=root_key), value=2).put()
+  def testGetLatestVersionNumberForSingularEntity(self):
+    root_key = ndb.Key('_SingularEntityRoot', 1)
+    _SingularEntity._GetRootModel()(key=root_key, current=1).put()
+    _SingularEntity(
+        key=ndb.Key('_SingularEntity', 1, parent=root_key), value=2).put()
 
-    self.assertEqual(1, _Entity.GetLatestVersionNumber())
+    self.assertEqual(1, _SingularEntity.GetLatestVersionNumber())
+
+  def testGetVersionForMultipleEntity(self):
+    entity = _MultipleEntity.Create('m1')
+    entity.value = 1
+    entity.Save()
+    entity.value = 2
+    entity.Save()
+    self.assertEqual(2, _MultipleEntity.GetVersion('m1').version)
+    self.assertEqual(2, _MultipleEntity.GetVersion('m1').value)
+    self.assertIsNone(_MultipleEntity.GetVersion('m1', version=0))
+    self.assertIsNone(_MultipleEntity.GetVersion('m1', version=3))
 
   def testGetLatestVersionNumberWhenNoRecordYet(self):
-    self.assertEqual(-1, _Entity.GetLatestVersionNumber())
+    self.assertEqual(-1, _SingularEntity.GetLatestVersionNumber())
+    self.assertEqual(-1, _MultipleEntity.GetLatestVersionNumber('m1'))
 
-  def testSaveNewVersion(self):
-    entity = _Entity()
+  def testSaveNewVersionForSingularEntity(self):
+    entity = _SingularEntity()
     entity.value = 1
-    key = entity.Save()
+    key, saved = entity.Save()
 
-    expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 1)
+    expected_key = ndb.Key('_SingularEntityRoot', 1, '_SingularEntity', 1)
     self.assertEqual(expected_key, key)
+    self.assertTrue(saved)
 
-    entity = _Entity.GetVersion()
+    entity = _SingularEntity.GetVersion()
     self.assertEqual(1, entity.version)
     self.assertEqual(1, entity.value)
 
-  def testSaveNewVersionAlreadyExist(self):
+  def testSaveNewVersionForMultipleEntity(self):
+    entity = _MultipleEntity.Create('m1')
+    entity.value = 1
+    key, saved = entity.Save()
+
+    expected_key = ndb.Key(
+        '_MultipleEntityRoot', 'm1', '_MultipleEntity', 1)
+    self.assertEqual(expected_key, key)
+    self.assertTrue(saved)
+
+  def testSaveNewVersionAlreadyExistsTryNextVersion(self):
     original_ndb_transaction = ndb.transaction
 
     def MockNdbTransaction(func, **options):
-      _Entity(key=ndb.Key('_EntityRoot', 1, '_Entity', 1), value=1).put()
-      _Entity(key=ndb.Key('_EntityRoot', 1, '_Entity', 2), value=2).put()
+      _SingularEntity(
+          key=ndb.Key(
+              '_SingularEntityRoot', 1, '_SingularEntity', 1), value=1).put()
+      _SingularEntity(
+          key=ndb.Key(
+              '_SingularEntityRoot', 1, '_SingularEntity', 2), value=2).put()
       return original_ndb_transaction(func, **options)
     self.mock(ndb, 'transaction', MockNdbTransaction)
 
-    entity = _Entity()
-    key = entity.Save()
-    expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 3)
+    entity = _SingularEntity()
+    key, saved = entity.Save()
+    expected_key = ndb.Key('_SingularEntityRoot', 1, '_SingularEntity', 3)
     self.assertEqual(expected_key, key)
+    self.assertTrue(saved)
+
+  def testSaveDataNewVersionAlreadyExistsNotTryNextVersion(self):
+    # When the call to SaveData() starts, another transaction had just beat this
+    # one to it. The expected result is the saved version from the first
+    # transaction.
+    original_ndb_transaction = ndb.transaction
+
+    def MockNdbTransaction(func, **options):
+      _SingularEntity(
+          key=ndb.Key(
+              '_SingularEntityRoot', 1, '_SingularEntity', 1), value=1).put()
+      return original_ndb_transaction(func, **options)
+    self.mock(ndb, 'transaction', MockNdbTransaction)
+
+    entity = _SingularEntity()
+    key, saved = entity.Save(should_try_next_version_number=False)
+    expected_key = ndb.Key('_SingularEntityRoot', 1, '_SingularEntity', 1)
+    self.assertEqual(expected_key, key)
+    self.assertFalse(saved)
+
+  def testSaveNewVersionAlreadyExistsNotTryNextVersion(self):
+    # When the call to Save() starts, another transaction had just beat this
+    # one to it. The expected result is the saved version from the first
+    # transaction should be returned.
+    original_save = VersionedModel.Save
+
+    def MockSave(*args, **kwargs):
+      _SingularEntity(
+          key=ndb.Key(
+              '_SingularEntityRoot', 1, '_SingularEntity', 1), value=1).put()
+      return original_save(*args, **kwargs)
+    self.mock(VersionedModel, 'Save', MockSave)
+
+    entity = _SingularEntity()
+    key, saved = entity.Save(should_try_next_version_number=False)
+    expected_key = ndb.Key('_SingularEntityRoot', 1, '_SingularEntity', 1)
+    self.assertEqual(expected_key, key)
+    self.assertFalse(saved)
 
   def testLikelyTransactionFailure(self):
     original_ndb_transaction = ndb.transaction
 
     calls = []
     def MockNdbTransaction(func, **options):
       if len(calls) < 1:
         calls.append(1)
         raise datastore_errors.Timeout()
       return original_ndb_transaction(func, **options)
     self.mock(ndb, 'transaction', MockNdbTransaction)
 
-    entity = _Entity()
-    key = entity.Save()
-    expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 1)
+    entity = _SingularEntity()
+    key, saved = entity.Save()
+    expected_key = ndb.Key('_SingularEntityRoot', 1, '_SingularEntity', 1)
     self.assertEqual(expected_key, key)
+    self.assertTrue(saved)
     self.assertEqual([1], calls)
 
   def testTransactionFailure(self):
     original_ndb_transaction = ndb.transaction
 
     calls = []
     def MockNdbTransaction(func, **options):
       if len(calls) < 1:
         calls.append(1)
         raise datastore_errors.BadRequestError()
       return original_ndb_transaction(func, **options)
     self.mock(ndb, 'transaction', MockNdbTransaction)
 
-    entity = _Entity()
-    key = entity.Save()
-    expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 1)
+    entity = _SingularEntity()
+    key, saved = entity.Save()
+    expected_key = ndb.Key('_SingularEntityRoot', 1, '_SingularEntity', 1)
     self.assertEqual(expected_key, key)
     self.assertEqual([1], calls)
+    self.assertTrue(saved)
+
+  def testRootStringId(self):
+    self.assertIsNone(_SingularEntity().root_string_id)
+    self.assertIsNone(
+        _SingularEntity(key=ndb.Key('_SingularEntity', 1)).root_string_id)
+    self.assertEqual('m1', _MultipleEntity.Create('m1').root_string_id)