Merge remote-tracking branch 'minjung/flushing_sqlalchemy_factory'

1 files changed, 32 insertions, 0 deletions

diff --git a/tests/test_alchemy.py b/tests/test_alchemy.py
index 9d7288a..5d8f275 100644
--- a/tests/test_alchemy.py
+++ b/tests/test_alchemy.py
@@ -23,6 +23,7 @@
 
 import factory
 from .compat import unittest
+import mock
 
 
 try:
@@ -55,6 +56,16 @@ class StandardFactory(SQLAlchemyModelFactory):
     foo = factory.Sequence(lambda n: 'foo%d' % n)
 
 
+class ForceFlushingStandardFactory(SQLAlchemyModelFactory):
+    class Meta:
+        model = models.StandardModel
+        sqlalchemy_session = mock.MagicMock()
+        force_flush = True
+
+    id = factory.Sequence(lambda n: n)
+    foo = factory.Sequence(lambda n: 'foo%d' % n)
+
+
 class NonIntegerPkFactory(SQLAlchemyModelFactory):
     class Meta:
         model = models.NonIntegerPk
@@ -103,6 +114,27 @@ class SQLAlchemyPkSequenceTestCase(unittest.TestCase):
 
 
 @unittest.skipIf(sqlalchemy is None, "SQLalchemy not installed.")
+class SQLAlchemyForceFlushTestCase(unittest.TestCase):
+    def setUp(self):
+        super(SQLAlchemyForceFlushTestCase, self).setUp()
+        ForceFlushingStandardFactory.reset_sequence(1)
+        ForceFlushingStandardFactory._meta.sqlalchemy_session.rollback()
+        ForceFlushingStandardFactory._meta.sqlalchemy_session.reset_mock()
+
+    def test_force_flush_called(self):
+        self.assertFalse(ForceFlushingStandardFactory._meta.sqlalchemy_session.flush.called)
+        ForceFlushingStandardFactory.create()
+        self.assertTrue(ForceFlushingStandardFactory._meta.sqlalchemy_session.flush.called)
+
+    def test_force_flush_not_called(self):
+        ForceFlushingStandardFactory._meta.force_flush = False
+        self.assertFalse(ForceFlushingStandardFactory._meta.sqlalchemy_session.flush.called)
+        ForceFlushingStandardFactory.create()
+        self.assertFalse(ForceFlushingStandardFactory._meta.sqlalchemy_session.flush.called)
+        ForceFlushingStandardFactory._meta.force_flush = True
+
+
+@unittest.skipIf(sqlalchemy is None, "SQLalchemy not installed.")
 class SQLAlchemyNonIntegerPkTestCase(unittest.TestCase):
     def setUp(self):
         super(SQLAlchemyNonIntegerPkTestCase, self).setUp()