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

appengine/findit/model/test/versioned_model_test.py

Index: appengine/findit/model/test/versioned_model_test.py
diff --git a/appengine/findit/model/test/versioned_model_test.py b/appengine/findit/model/test/versioned_model_test.py
index cfe40a964c9d464f3560c13d179b607a3ce89958..efc7f727f6a516b042c410723c590af2a4a00d8c 100644
--- a/appengine/findit/model/test/versioned_model_test.py
+++ b/appengine/findit/model/test/versioned_model_test.py
@@ -23,58 +23,100 @@ class VersionedModelTest(testing.AppengineTestCase):
     self.assertEqual(3, root_model_class(current=3).current)
 
   def testDefaultVersionIsZero(self):
-    entity = _Entity()
-    self.assertEqual(0, entity.version)
+    self.assertEqual(0, _Entity.Create().version)
+    self.assertEqual(0, _Entity.Create('m1').version)
+
+  def testRootId(self):
+    self.assertEqual(1, _Entity.Create()._root_id)
+    self.assertEqual(
+        1, _Entity(key=ndb.Key('_EntityRoot', 1, '_Entity', 2))._root_id)
+    self.assertEqual('m1', _Entity.Create('m1')._root_id)
+
+  def testCreate(self):
+    self.assertEqual(_Entity(key=ndb.Key('_EntityRoot', 'm1', '_Entity', 0)),
+                     _Entity.Create('m1'))
 
   def testGetMostRecentVersionWhenNoData(self):
-    entity = _Entity.GetVersion()
-    self.assertIsNone(entity)
+    self.assertIsNone(_Entity.GetVersion())
+    self.assertIsNone(_Entity.GetVersion('m1'))
 
   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()
+    entity = _Entity.Create()
+    entity.value = 1
+    entity.Save()
+    entity.value = 2
+    entity.Save()
 
     entity = _Entity.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 testGetSpecificVersion(self):
+    entity = _Entity.Create()
+    entity.value = 1
+    entity.Save()
+    entity.value = 2
+    entity.Save()
 
-    entity = _Entity.GetVersion(2)
-    self.assertEqual(2, entity.version)
-    self.assertEqual(2, entity.value)
-    self.assertIsNone(_Entity.GetVersion(0))
-    self.assertIsNone(_Entity.GetVersion(3))
+    entity_version_2 = _Entity.GetVersion(version=2)
+    self.assertEqual(2, entity_version_2.version)
+    self.assertEqual(2, entity_version_2.value)
+    self.assertIsNone(_Entity.GetVersion(version=0))
+    self.assertIsNone(_Entity.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()
-
+    entity = _Entity.Create()
+    entity.Save()
     self.assertEqual(1, _Entity.GetLatestVersionNumber())
 
+  def testGetVersionForEntityWithKeyId(self):
+    entity = _Entity.Create('m1')
+    entity.value = 1
+    entity.Save()
+    entity.value = 2
+    entity.Save()
+
+    self.assertEqual(2, _Entity.GetVersion('m1').version)
+    self.assertEqual(2, _Entity.GetVersion('m1').value)
+    self.assertIsNone(_Entity.GetVersion('m1', version=0))
+    self.assertIsNone(_Entity.GetVersion('m1', version=3))
+
   def testGetLatestVersionNumberWhenNoRecordYet(self):
     self.assertEqual(-1, _Entity.GetLatestVersionNumber())
+    self.assertEqual(-1, _Entity.GetLatestVersionNumber('m1'))
 
   def testSaveNewVersion(self):
-    entity = _Entity()
+    entity = _Entity.Create()
     entity.value = 1
-    key = entity.Save()
+    key, saved = entity.Save()
 
     expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 1)
     self.assertEqual(expected_key, key)
+    self.assertTrue(saved)
 
     entity = _Entity.GetVersion()
     self.assertEqual(1, entity.version)
     self.assertEqual(1, entity.value)
 
-  def testSaveNewVersionAlreadyExist(self):
+  def testSaveEntityWithRootId(self):
+    entity = _Entity.Create('m1')
+    entity.value = 1
+    entity.Save()
+    entity.value = 3
+    entity.Save()
+
+    self.assertEqual(entity.version, 2)
+    self.assertEqual(entity.value, 3)
+
+  def testSaveNewVersionOfEntityWithKeyId(self):
+    entity = _Entity.Create('m1')
+    entity.value = 1
+    key, saved = entity.Save()
+    expected_key = ndb.Key('_EntityRoot', 'm1', '_Entity', 1)
+    self.assertEqual(expected_key, key)
+    self.assertTrue(saved)
+
+  def testSaveNewVersionWithConflictRetry(self):
     original_ndb_transaction = ndb.transaction
 
     def MockNdbTransaction(func, **options):
@@ -83,10 +125,43 @@ class VersionedModelTest(testing.AppengineTestCase):
       return original_ndb_transaction(func, **options)
     self.mock(ndb, 'transaction', MockNdbTransaction)
 
-    entity = _Entity()
-    key = entity.Save()
+    entity = _Entity.Create()
+    key, saved = entity.Save()
     expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 3)
     self.assertEqual(expected_key, key)
+    self.assertTrue(saved)
+
+  def testSaveNewVersionWithConflictBeforeTransactionNoRetry(self):
+    original_ndb_transaction = ndb.transaction
+
+    def MockNdbTransaction(func, **options):
+      # To simulate another transaction occuring before this one, write a new
+      # version first before the transaction is called.
+      _Entity(key=ndb.Key('_EntityRoot', 1, '_Entity', 1), value=1).put()
+      return original_ndb_transaction(func, **options)
+    self.mock(ndb, 'transaction', MockNdbTransaction)
+
+    entity = _Entity.Create()
+    key, saved = entity.Save(retry_on_conflict=False)
+    expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 1)
+    self.assertEqual(expected_key, key)
+    self.assertFalse(saved)
+
+  def testSaveNewVersionConflictBeforeSaving(self):
+    original_save = VersionedModel.Save
+
+    def MockSave(*args, **kwargs):
+      # Smulate another transaction beating this call to Save() to creating a
+      # new entity.
+      _Entity(key=ndb.Key('_EntityRoot', 1, '_Entity', 1), value=1).put()
+      return original_save(*args, **kwargs)
+    self.mock(VersionedModel, 'Save', MockSave)
+
+    entity = _Entity.Create()
+    key, saved = entity.Save(retry_on_conflict=False)
+    expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 1)
+    self.assertEqual(expected_key, key)
+    self.assertFalse(saved)
 
   def testLikelyTransactionFailure(self):
     original_ndb_transaction = ndb.transaction
@@ -99,10 +174,11 @@ class VersionedModelTest(testing.AppengineTestCase):
       return original_ndb_transaction(func, **options)
     self.mock(ndb, 'transaction', MockNdbTransaction)
 
-    entity = _Entity()
-    key = entity.Save()
+    entity = _Entity.Create()
+    key, saved = entity.Save()
     expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 1)
     self.assertEqual(expected_key, key)
+    self.assertTrue(saved)
     self.assertEqual([1], calls)
 
   def testTransactionFailure(self):
@@ -116,8 +192,9 @@ class VersionedModelTest(testing.AppengineTestCase):
       return original_ndb_transaction(func, **options)
     self.mock(ndb, 'transaction', MockNdbTransaction)
 
-    entity = _Entity()
-    key = entity.Save()
+    entity = _Entity.Create()
+    key, saved = entity.Save()
     expected_key = ndb.Key('_EntityRoot', 1, '_Entity', 1)
     self.assertEqual(expected_key, key)
     self.assertEqual([1], calls)
+    self.assertTrue(saved)