testkit_core/handlers/
with_transaction.rs

1// src/handlers/with_transaction.rs
2use async_trait::async_trait;
3use std::fmt::Debug;
4use std::marker::PhantomData;
5
6use crate::{
7    DatabaseBackend, TestContext,
8    handlers::{IntoTransactionHandler, TransactionHandler},
9};
10
11/// Handler for executing functions within a transaction
12pub struct TransactionFnHandler<DB, F>
13where
14    DB: DatabaseBackend + Send + Sync + Debug + 'static,
15    F: Send + Sync + 'static,
16{
17    transaction_fn: F,
18    _phantom: PhantomData<DB>,
19}
20
21/// Create a new transaction function handler
22pub fn with_transaction<DB, F, Fut>(transaction_fn: F) -> TransactionFnHandler<DB, F>
23where
24    DB: DatabaseBackend + Send + Sync + Debug + 'static,
25    Fut: std::future::Future<Output = Result<(), DB::Error>> + Send + 'static,
26    F: FnOnce(&mut <DB as DatabaseBackend>::Connection) -> Fut + Send + Sync + 'static,
27{
28    TransactionFnHandler {
29        transaction_fn,
30        _phantom: PhantomData,
31    }
32}
33
34#[async_trait]
35impl<DB, F, Fut> TransactionHandler<DB> for TransactionFnHandler<DB, F>
36where
37    DB: DatabaseBackend + Send + Sync + Debug + 'static,
38    Fut: std::future::Future<Output = Result<(), DB::Error>> + Send + 'static,
39    F: FnOnce(&mut <DB as DatabaseBackend>::Connection) -> Fut + Send + Sync + 'static,
40{
41    type Item = ();
42    type Error = DB::Error;
43
44    async fn execute(self, ctx: &mut TestContext<DB>) -> Result<Self::Item, Self::Error> {
45        // Acquire a connection from the database
46        let mut conn = ctx.db.acquire_connection().await?;
47
48        // Execute the transaction function
49        let result = (self.transaction_fn)(&mut conn).await;
50
51        // Release the connection
52        ctx.db.release_connection(conn).await?;
53
54        result
55    }
56}
57
58impl<DB, F, Fut> IntoTransactionHandler<DB> for TransactionFnHandler<DB, F>
59where
60    DB: DatabaseBackend + Send + Sync + Debug + 'static,
61    Fut: std::future::Future<Output = Result<(), DB::Error>> + Send + 'static,
62    F: FnOnce(&mut <DB as DatabaseBackend>::Connection) -> Fut + Send + Sync + 'static,
63{
64    type Handler = Self;
65    type Item = ();
66    type Error = DB::Error;
67
68    fn into_transaction_handler(self) -> Self::Handler {
69        self
70    }
71}
72
73/// Handler that combines a database instance with transaction management
74pub struct DatabaseTransactionHandler<DB, F>
75where
76    DB: DatabaseBackend + Send + Sync + Debug + 'static,
77    F: Send + Sync + 'static,
78{
79    /// The database instance
80    db: crate::TestDatabaseInstance<DB>,
81    /// The transaction function
82    transaction_fn: F,
83}
84
85/// Create a new database transaction handler
86pub fn with_db_transaction<DB, F, Fut>(
87    db: crate::TestDatabaseInstance<DB>,
88    transaction_fn: F,
89) -> DatabaseTransactionHandler<DB, F>
90where
91    DB: DatabaseBackend + Send + Sync + Debug + 'static,
92    Fut: std::future::Future<Output = Result<(), DB::Error>> + Send + 'static,
93    F: FnOnce(&mut <DB as DatabaseBackend>::Connection) -> Fut + Send + Sync + 'static,
94{
95    DatabaseTransactionHandler { db, transaction_fn }
96}
97
98#[async_trait]
99impl<DB, F, Fut> TransactionHandler<DB> for DatabaseTransactionHandler<DB, F>
100where
101    DB: DatabaseBackend + Send + Sync + Debug + 'static,
102    Fut: std::future::Future<Output = Result<(), DB::Error>> + Send + 'static,
103    F: FnOnce(&mut <DB as DatabaseBackend>::Connection) -> Fut + Send + Sync + 'static,
104{
105    type Item = TestContext<DB>;
106    type Error = DB::Error;
107
108    async fn execute(self, _ctx: &mut TestContext<DB>) -> Result<Self::Item, Self::Error> {
109        // Create a new context with the database
110        let ctx = TestContext::new(self.db);
111
112        // Acquire a connection
113        let mut conn = ctx.db.acquire_connection().await?;
114
115        // Execute the transaction function
116        (self.transaction_fn)(&mut conn).await?;
117
118        // Release the connection
119        ctx.db.release_connection(conn).await?;
120
121        Ok(ctx)
122    }
123}
124
125impl<DB, F, Fut> IntoTransactionHandler<DB> for DatabaseTransactionHandler<DB, F>
126where
127    DB: DatabaseBackend + Send + Sync + Debug + 'static,
128    Fut: std::future::Future<Output = Result<(), DB::Error>> + Send + 'static,
129    F: FnOnce(&mut <DB as DatabaseBackend>::Connection) -> Fut + Send + Sync + 'static,
130{
131    type Handler = Self;
132    type Item = TestContext<DB>;
133    type Error = DB::Error;
134
135    fn into_transaction_handler(self) -> Self::Handler {
136        self
137    }
138}