Fix: Support transactions in endpoints with test tool (serverpod#2791)

packages/serverpod_test/lib/serverpod_test_public_exports.dart

@@ -0,0 +1,17 @@
+library serverpod_test;
+
+// ATTENTION: This file should never be imported directly.
+// It is used to re-export the public parts of the test tools
+// in the generated test tools file.
+export 'serverpod_test.dart'
+    show
+        AuthenticationOverride,
+        ConnectionClosedException,
+        flushMicrotasks,
+        InvalidConfigurationException,
+        ResetTestSessions,
+        RollbackDatabase,
+        ServerpodInsufficientAccessException,
+        ServerpodUnauthenticatedException,
+        TestClosure,
+        TestSession;

packages/serverpod_test/lib/src/test_database_proxy.dart

@@ -0,0 +1,234 @@
+import 'package:serverpod/serverpod.dart';
+import 'package:serverpod_test/src/transaction_manager.dart';
+
+import 'with_serverpod.dart';
+
+/// A database proxy that forwards all calls to the provided database.
+class TestDatabaseProxy implements Database {
+  final Database _db;
+  final RollbackDatabase _rollbackDatabase;
+  final TransactionManager _transactionManager;
+
+  /// Creates a new [TestDatabaseProxy]
+  TestDatabaseProxy(this._db, this._rollbackDatabase, this._transactionManager);
+
+  @override
+  Future<int> count<T extends TableRow>({
+    Expression? where,
+    int? limit,
+    bool useCache = true,
+    Transaction? transaction,
+  }) {
+    return _db.count<T>(
+      where: where,
+      limit: limit,
+      useCache: useCache,
+      transaction: transaction,
+    );
+  }
+
+  @override
+  Future<List<T>> delete<T extends TableRow>(
+    List<T> rows, {
+    Transaction? transaction,
+  }) {
+    return _db.delete<T>(rows, transaction: transaction);
+  }
+
+  @override
+  Future<T> deleteRow<T extends TableRow>(
+    T row, {
+    Transaction? transaction,
+  }) {
+    return _db.deleteRow<T>(row, transaction: transaction);
+  }
+
+  @override
+  Future<List<T>> deleteWhere<T extends TableRow>({
+    required Expression where,
+    Transaction? transaction,
+  }) {
+    return _db.deleteWhere<T>(where: where, transaction: transaction);
+  }
+
+  @override
+  Future<List<T>> find<T extends TableRow>({
+    Expression? where,
+    int? limit,
+    int? offset,
+    Column? orderBy,
+    List<Order>? orderByList,
+    bool orderDescending = false,
+    Transaction? transaction,
+    Include? include,
+  }) {
+    return _db.find<T>(
+      where: where,
+      limit: limit,
+      offset: offset,
+      orderBy: orderBy,
+      orderByList: orderByList,
+      orderDescending: orderDescending,
+      transaction: transaction,
+      include: include,
+    );
+  }
+
+  @override
+  Future<T?> findById<T extends TableRow>(
+    int id, {
+    Transaction? transaction,
+    Include? include,
+  }) {
+    return _db.findById<T>(id, transaction: transaction, include: include);
+  }
+
+  @override
+  Future<T?> findFirstRow<T extends TableRow>({
+    Expression? where,
+    int? offset,
+    Column? orderBy,
+    List<Order>? orderByList,
+    bool orderDescending = false,
+    Transaction? transaction,
+    Include? include,
+  }) {
+    return _db.findFirstRow<T>(
+      where: where,
+      offset: offset,
+      orderBy: orderBy,
+      orderByList: orderByList,
+      orderDescending: orderDescending,
+      transaction: transaction,
+      include: include,
+    );
+  }
+
+  @override
+  Future<List<T>> insert<T extends TableRow>(
+    List<T> rows, {
+    Transaction? transaction,
+  }) {
+    return _db.insert<T>(rows, transaction: transaction);
+  }
+
+  @override
+  Future<T> insertRow<T extends TableRow>(
+    T row, {
+    Transaction? transaction,
+  }) {
+    return _db.insertRow<T>(row, transaction: transaction);
+  }
+
+  @override
+  Future<bool> testConnection() {
+    return _db.testConnection();
+  }
+
+  @override
+  Future<R> transaction<R>(
+    TransactionFunction<R> transactionFunction, {
+    isUserCall = true,
+  }) async {
+    if (!isUserCall || _rollbackDatabase == RollbackDatabase.disabled) {
+      return _db.transaction(transactionFunction);
+    }
+
+    var localTransaction = _transactionManager.currentTransaction;
+    if (localTransaction == null) {
+      throw StateError('No ongoing transaction.');
+    }
+
+    try {
+      await _transactionManager.addSavePoint(lock: true);
+    } on ConcurrentTransactionsException {
+      throw InvalidConfigurationException(
+        'Concurrent calls to transaction are not supported when database rollbacks are enabled. '
+        'Disable rolling back the database by setting `rollbackDatabase` to `RollbackDatabase.disabled`.',
+      );
+    }
+
+    try {
+      var result = await transactionFunction(localTransaction);
+      await _transactionManager.removePreviousSavePoint(unlock: true);
+      return result;
+    } catch (e) {
+      await _transactionManager.rollbacktoPreviousSavePoint(unlock: true);
+      rethrow;
+    }
+  }
+
+  @override
+  Future<int> unsafeExecute(
+    String query, {
+    int? timeoutInSeconds,
+    Transaction? transaction,
+    QueryParameters? parameters,
+  }) {
+    return _db.unsafeExecute(
+      query,
+      timeoutInSeconds: timeoutInSeconds,
+      transaction: transaction,
+      parameters: parameters,
+    );
+  }
+
+  @override
+  Future<DatabaseResult> unsafeQuery(
+    String query, {
+    int? timeoutInSeconds,
+    Transaction? transaction,
+    QueryParameters? parameters,
+  }) {
+    return _db.unsafeQuery(
+      query,
+      timeoutInSeconds: timeoutInSeconds,
+      transaction: transaction,
+      parameters: parameters,
+    );
+  }
+
+  @override
+  Future<int> unsafeSimpleExecute(
+    String query, {
+    int? timeoutInSeconds,
+    Transaction? transaction,
+  }) {
+    return _db.unsafeSimpleExecute(
+      query,
+      timeoutInSeconds: timeoutInSeconds,
+      transaction: transaction,
+    );
+  }
+
+  @override
+  Future<DatabaseResult> unsafeSimpleQuery(
+    String query, {
+    int? timeoutInSeconds,
+    Transaction? transaction,
+  }) {
+    return _db.unsafeSimpleQuery(
+      query,
+      timeoutInSeconds: timeoutInSeconds,
+      transaction: transaction,
+    );
+  }
+
+  @override
+  Future<List<T>> update<T extends TableRow>(
+    List<T> rows, {
+    List<Column>? columns,
+    Transaction? transaction,
+  }) {
+    return _db.update<T>(rows, columns: columns, transaction: transaction);
+  }
+
+  @override
+  Future<T> updateRow<T extends TableRow>(
+    T row, {
+    List<Column>? columns,
+    Transaction? transaction,
+  }) {
+    return _db.updateRow<T>(row, columns: columns, transaction: transaction);
+  }
+}

packages/serverpod_test/lib/src/test_serverpod.dart

@@ -1,4 +1,6 @@
 import 'package:serverpod/serverpod.dart';
+import 'package:serverpod_test/src/test_database_proxy.dart';
+import 'package:serverpod_test/src/transaction_manager.dart';
 import 'package:serverpod_test/src/with_serverpod.dart';
 
 /// Internal test endpoints interface that contains implementation details
@@ -16,16 +18,35 @@ abstract interface class InternalTestEndpoints {
 class InternalServerpodSession extends Session {
   /// The transaction that is used by the session.
   @override
-  Transaction? transaction;
+  Transaction? get transaction => transactionManager.currentTransaction;
+
+  @override
+  TestDatabaseProxy get db => _dbProxy;
+
+  late TestDatabaseProxy _dbProxy;
+
+  /// The database test configuration.
+  final RollbackDatabase rollbackDatabase;
+
+  /// The transaction manager to manage the Serverpod session's transactions.
+  late final TransactionManager transactionManager;
 
   /// Creates a new internal serverpod session.
   InternalServerpodSession({
     required super.endpoint,
     required super.method,
     required super.server,
     required super.enableLogging,
-    this.transaction,
-  });
+    required this.rollbackDatabase,
+    TransactionManager? transactionManager,
+  }) {
+    this.transactionManager = transactionManager ?? TransactionManager(this);
+    _dbProxy = TestDatabaseProxy(
+      super.db,
+      rollbackDatabase,
+      this.transactionManager,
+    );
+  }
 }
 
 List<String> _getServerpodStartUpArgs(String? runMode, bool? applyMigrations) =>
@@ -84,16 +105,18 @@ class TestServerpod<T extends InternalTestEndpoints> {
   /// Creates a new Serverpod session.
   InternalServerpodSession createSession({
     bool enableLogging = false,
-    Transaction? transaction,
+    required RollbackDatabase rollbackDatabase,
     String endpoint = '',
     String method = '',
+    TransactionManager? transactionManager,
   }) {
     return InternalServerpodSession(
       server: _serverpod.server,
       enableLogging: enableLogging,
       endpoint: endpoint,
       method: method,
-      transaction: transaction,
+      rollbackDatabase: rollbackDatabase,
+      transactionManager: transactionManager,
     );
   }
 }

packages/serverpod_test/lib/src/test_session.dart

@@ -70,29 +70,22 @@ class InternalTestSession extends TestSession {
   @override
   MessageCentralAccess get messages => serverpodSession.messages;
 
-  Transaction? _transaction;
   @override
-  Transaction? get transaction => _transaction;
-
-  set transaction(Transaction? transaction) {
-    _transaction = transaction;
-    serverpodSession.transaction = transaction;
-  }
+  Transaction? get transaction =>
+      serverpodSession.transactionManager.currentTransaction;
 
   /// Creates a new internal test session.
   InternalTestSession(
     TestServerpod testServerpod, {
     AuthenticationOverride? authenticationOverride,
     InternalTestSession? sessionWithDatabaseConnection,
-    Transaction? transaction,
     required bool enableLogging,
     required List<InternalTestSession> allTestSessions,
     required this.serverpodSession,
   })  : _allTestSessions = allTestSessions,
         _authenticationOverride = authenticationOverride,
         _testServerpod = testServerpod,
-        _enableLogging = enableLogging,
-        _transaction = transaction {
+        _enableLogging = enableLogging {
     _allTestSessions.add(this);
     _configureServerpodSession(serverpodSession);
   }
@@ -106,17 +99,17 @@ class InternalTestSession extends TestSession {
   }) {
     var newServerpodSession = _testServerpod.createSession(
       enableLogging: enableLogging ?? _enableLogging,
-      transaction: transaction,
       endpoint: endpoint,
       method: method,
+      rollbackDatabase: serverpodSession.rollbackDatabase,
+      transactionManager: serverpodSession.transactionManager,
     );
 
     return InternalTestSession(
       _testServerpod,
       allTestSessions: _allTestSessions,
       authenticationOverride: authentication ?? _authenticationOverride,
       enableLogging: enableLogging ?? _enableLogging,
-      transaction: transaction,
       serverpodSession: newServerpodSession,
     );
   }
@@ -134,8 +127,9 @@ class InternalTestSession extends TestSession {
     await serverpodSession.close();
     _authenticationOverride = null;
     serverpodSession = _testServerpod.createSession(
-      transaction: _transaction,
       enableLogging: _enableLogging,
+      rollbackDatabase: serverpodSession.rollbackDatabase,
+      transactionManager: serverpodSession.transactionManager,
     );
     _configureServerpodSession(serverpodSession);
   }