From 9cc226746bb1f2d948fb84e0b9a918178dd7a2f4 Mon Sep 17 00:00:00 2001
From: gauravr <gaurav@stck.me>
Date: Tue, 18 Feb 2025 07:00:31 +0530
Subject: [PATCH] Cleaning

---
 tests/test_piccolo.py | 27 +++++++++++++--------------
 1 file changed, 13 insertions(+), 14 deletions(-)

diff --git a/tests/test_piccolo.py b/tests/test_piccolo.py
index fb82a05..0c58968 100644
--- a/tests/test_piccolo.py
+++ b/tests/test_piccolo.py
@@ -2,11 +2,12 @@
 
 from piccolo import columns as col
 from piccolo.engine.postgres import PostgresEngine
+import pytest
 
 import settings
 from apphelpers.db.piccolo import (
     BaseTable,
-    dbtransaction_ctx,
+    dbtransaction,
     destroy_db_from_basetable,
     setup_db_from_basetable,
 )
@@ -19,6 +20,7 @@
         password=settings.DB_PASS,
     )
 )
+dbtransaction_wrapper = dbtransaction(db)
 
 
 class Book(BaseTable, db=db):
@@ -43,25 +45,22 @@ def teardown_module():
 
 
 async def add_with_tr():
-    async with dbtransaction_ctx(db):
-        name = "The Pillars of the Earth"
-        await _add_book(name)
+    add_book = dbtransaction_wrapper(_add_book)
+    name = "The Pillars of the Earth"
+    await add_book(name)
     names = [b[Book.name] for b in await Book.select().run()]
     assert name in names
 
-    try:
-        async with dbtransaction_ctx(db):
-            name = "The Cathedral and the Bazaar"
-            await _add_book_loser(name)
-    except NameError:
-        pass
-
+    add_book_loser = dbtransaction_wrapper(_add_book_loser)
+    name = "The Cathedral and the Bazaar"
+    with pytest.raises(NameError):
+        await add_book_loser(name)
     names = [b[Book.name] for b in await Book.select().run()]
     assert name not in names
 
-    async with dbtransaction_ctx(db):
-        name = "The Ego Trick"
-        await _add_book(name)
+    add_book = dbtransaction_wrapper(_add_book)
+    name = "The Ego Trick"
+    await add_book(name)
     names = [b[Book.name] for b in await Book.select().run()]
     assert name in names