Use newer APIs.

coleifer · Sep 22, 2015 · db8a39f · db8a39f
1 parent e8bedc0
commit db8a39f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 8 deletions.
diff --git a/huey/peewee_helpers/__init__.py b/huey/peewee_helpers/__init__.py
@@ -4,12 +4,9 @@
 def _transaction(db, fn):
     @wraps(fn)
     def inner(*args, **kwargs):
-        db.get_conn()
-        try:
-            with db.transaction():
-                return fn(*args, **kwargs)
-        finally:
-            db.close()
+        # Execute function in its own connection, in a transaction.
+        with db.execution_context(with_transaction=True):
+            return fn(*args, **kwargs)
     return inner
 
 def db_task(huey, db, *args, **kwargs):

diff --git a/huey/tests/peewee_tests.py b/huey/tests/peewee_tests.py
@@ -24,7 +24,7 @@ def inner(*args, **kwargs):
             return fn(*args, **kwargs)
         return inner
     connect = record_call(SqliteDatabase.connect)
-    close = record_call(SqliteDatabase.close)
+    _close = record_call(SqliteDatabase._close)
     transaction = record_call(SqliteDatabase.transaction)
 
 db = MockSqliteDatabase('test.huey.db')
@@ -58,5 +58,5 @@ def test_helper(self):
         test_db_task('foo')
         self.assertEqual(STATE, ['connect'])
         huey.execute(huey.dequeue())
-        self.assertEqual(STATE, ['connect', 'transaction', 'create', 'close'])
+        self.assertEqual(STATE, ['connect', 'transaction', 'create', '_close'])
         self.assertEqual(Value.select().count(), 1)