Allow nested atomic blocks, fix tests

flask_signalbus/atomic.py

@@ -113,8 +113,8 @@ def f():
         def wrapper(*args, **kwargs):
             session = self.session
             session_info = session.info
-            assert not session_info.get(_ATOMIC_FLAG_SESSION_INFO_KEY), \
-                '"atomic" blocks can not be nested'
+            if session_info.get(_ATOMIC_FLAG_SESSION_INFO_KEY):
+                return func(*args, **kwargs)
             f = retry_on_deadlock(session)(func)
             session_info[_ATOMIC_FLAG_SESSION_INFO_KEY] = True
             try:

tests/test_atomic.py

@@ -26,33 +26,46 @@ def f(x):
 def test_execute_atomic(atomic_db):
     db = atomic_db
     commit = Mock()
+    rollback = Mock()
     db.session.commit = commit
+    db.session.rollback = rollback
     var = 1
 
     with pytest.raises(RuntimeError):
         @db.execute_atomic
         def f1():
             raise RuntimeError
     commit.assert_not_called()
-
-    with pytest.raises(AssertionError):
-        @db.execute_atomic
-        def f2():
-            @db.execute_atomic
-            def recursive():
-                pass
-    commit.assert_not_called()
+    rollback.assert_called_once()
 
     @db.execute_atomic
-    def f3():
+    def f2():
         assert var == 1
         return 666
     commit.assert_called_once()
-    assert f3 == 666
+    rollback.assert_called_once()
+    assert f2 == 666
 
     assert db.execute_atomic(lambda x: x, 777) == 777
 
 
+def test_nested_execute_atomic(atomic_db):
+    db = atomic_db
+    commit = Mock()
+    rollback = Mock()
+    db.session.commit = commit
+    db.session.rollback = rollback
+
+    @db.execute_atomic
+    def f1():
+        @db.execute_atomic
+        def f2():
+            pass
+
+    commit.assert_called_once()
+    rollback.assert_not_called()
+
+
 def test_retry_on_integrity_error(atomic_db, AtomicModel):
     db = atomic_db
     o = AtomicModel(