ITestDbContextProvider is IAsyncDisposable

Author: PawelGerr

src/Thinktecture.EntityFrameworkCore.SqlServer.Testing/EntityFrameworkCore/Testing/ITestIsolationOptions.cs

@@ -10,21 +10,43 @@ namespace Thinktecture.EntityFrameworkCore.Testing;
 /// </summary>
 public interface ITestIsolationOptions
 {
+   /// <summary>
+   /// No test isolation, no cleanup.
+   /// </summary>
+   public static readonly ITestIsolationOptions None = new NoCleanup();
+
    /// <summary>
    /// Test isolation via ambient transaction.
    /// </summary>
    public static readonly ITestIsolationOptions SharedTablesAmbientTransaction = new ShareTablesIsolation();
 
    /// <summary>
-   /// Rollbacks migrations and then deletes database objects (like tables).
+   /// Rollbacks migrations and then deletes database objects (like tables) with a schema used by the tests.
    /// </summary>
    public static readonly ITestIsolationOptions RollbackMigrationsAndCleanup = new RollbackMigrationsAndCleanupDatabase();
 
    /// <summary>
-   /// Deletes database objects (like tables).
+   /// Deletes database objects (like tables) with a schema used by the tests.
    /// </summary>
    public static readonly ITestIsolationOptions CleanupOnly = new CleanupDatabase();
 
+   /// <summary>
+   /// Deletes all records from the tables.
+   /// </summary>
+   public static readonly ITestIsolationOptions TruncateTables = new TruncateAllTables();
+
+   /// <summary>
+   /// Performs custom cleanup.
+   /// </summary>
+   /// <param name="cleanup">Callback that performs the actual cleanup.</param>
+   /// <typeparam name="T">Type of the <see cref="DbContext"/></typeparam>
+   /// <returns></returns>
+   public static ITestIsolationOptions Custom<T>(Func<T, string, CancellationToken, Task> cleanup)
+      where T : DbContext
+   {
+      return new CustomCleanup<T>(cleanup);
+   }
+
    /// <summary>
    /// Indicator, whether the database needs cleanup.
    /// </summary>
@@ -33,60 +55,103 @@ public interface ITestIsolationOptions
    /// <summary>
    /// Cleanup of the database.
    /// </summary>
-   void Cleanup(DbContext dbContext, string schema);
+   ValueTask CleanupAsync(DbContext dbContext, string schema, CancellationToken cancellationToken);
+
+   private class NoCleanup : ITestIsolationOptions
+   {
+      public bool NeedsCleanup => false;
+
+      public ValueTask CleanupAsync(DbContext dbContext, string schema, CancellationToken cancellationToken)
+      {
+         return ValueTask.CompletedTask;
+      }
+   }
 
    private class ShareTablesIsolation : ITestIsolationOptions
    {
       public bool NeedsCleanup => false;
 
-      public void Cleanup(DbContext dbContext, string schema)
+      public ValueTask CleanupAsync(DbContext dbContext, string schema, CancellationToken cancellationToken)
       {
+         return ValueTask.CompletedTask;
       }
    }
 
    private class RollbackMigrationsAndCleanupDatabase : ITestIsolationOptions
    {
       public bool NeedsCleanup => true;
 
-      public void Cleanup(DbContext dbContext, string schema)
+      public async ValueTask CleanupAsync(DbContext dbContext, string schema, CancellationToken cancellationToken)
       {
-         RollbackMigrations(dbContext);
-         DeleteDatabaseObjects(dbContext, schema);
+         await RollbackMigrationsAsync(dbContext, cancellationToken);
+         await DeleteDatabaseObjectsAsync(dbContext, schema, cancellationToken);
       }
    }
 
    private class CleanupDatabase : ITestIsolationOptions
    {
       public bool NeedsCleanup => true;
 
-      public void Cleanup(DbContext dbContext, string schema)
+      public async ValueTask CleanupAsync(DbContext dbContext, string schema, CancellationToken cancellationToken)
+      {
+         await DeleteDatabaseObjectsAsync(dbContext, schema, cancellationToken);
+      }
+   }
+
+   private class CustomCleanup<T> : ITestIsolationOptions
+      where T : DbContext
+   {
+      private readonly Func<T, string, CancellationToken, Task> _cleanup;
+
+      public bool NeedsCleanup => true;
+
+      public CustomCleanup(Func<T, string, CancellationToken, Task> cleanup)
+      {
+         _cleanup = cleanup;
+      }
+
+      public async ValueTask CleanupAsync(DbContext dbContext, string schema, CancellationToken cancellationToken)
+      {
+         await _cleanup((T)dbContext, schema, cancellationToken);
+      }
+   }
+
+   private class TruncateAllTables : ITestIsolationOptions
+   {
+      public bool NeedsCleanup => true;
+
+      public async ValueTask CleanupAsync(DbContext dbContext, string schema, CancellationToken cancellationToken)
       {
-         DeleteDatabaseObjects(dbContext, schema);
+         foreach (var entityType in dbContext.Model.GetEntityTypes())
+         {
+            if (entityType.GetTableName() is not null)
+               await dbContext.TruncateTableAsync(entityType.ClrType, cancellationToken);
+         }
       }
    }
 
    /// <summary>
    /// Rollbacks all migrations.
    /// </summary>
-   protected static void RollbackMigrations(DbContext dbContext)
+   protected static async Task RollbackMigrationsAsync(DbContext dbContext, CancellationToken cancellationToken)
    {
       ArgumentNullException.ThrowIfNull(dbContext);
 
-      dbContext.GetService<IMigrator>().Migrate("0");
+      await dbContext.GetService<IMigrator>().MigrateAsync("0", cancellationToken);
    }
 
    /// <summary>
    /// Deletes database objects (like tables) with provided <paramref name="schema"/>.
    /// </summary>
-   protected static void DeleteDatabaseObjects(DbContext ctx, string schema)
+   protected static async Task DeleteDatabaseObjectsAsync(DbContext ctx, string schema, CancellationToken cancellationToken)
    {
       ArgumentNullException.ThrowIfNull(ctx);
       ArgumentNullException.ThrowIfNull(schema);
 
       var sqlHelper = ctx.GetService<ISqlGenerationHelper>();
 
-      ctx.Database.ExecuteSqlRaw(GetSqlForCleanup(), new SqlParameter("@schema", schema));
-      ctx.Database.ExecuteSqlRaw(GetDropSchemaSql(sqlHelper, schema));
+      await ctx.Database.ExecuteSqlRawAsync(GetSqlForCleanup(), new object[] { new SqlParameter("@schema", schema) }, cancellationToken);
+      await ctx.Database.ExecuteSqlRawAsync(GetDropSchemaSql(sqlHelper, schema), cancellationToken);
    }
 
    private static string GetSqlForCleanup()

src/Thinktecture.EntityFrameworkCore.SqlServer.Testing/EntityFrameworkCore/Testing/SqlServerDbContextIntegrationTests.cs

@@ -1,3 +1,4 @@
+using Xunit;
 using Xunit.Abstractions;
 
 namespace Thinktecture.EntityFrameworkCore.Testing;
@@ -6,7 +7,7 @@ namespace Thinktecture.EntityFrameworkCore.Testing;
 /// A base class for integration tests using EF Core along with SQL Server.
 /// </summary>
 /// <typeparam name="T">Type of the database context.</typeparam>
-public abstract class SqlServerDbContextIntegrationTests<T> : ITestDbContextProvider<T>
+public abstract class SqlServerDbContextIntegrationTests<T> : ITestDbContextProvider<T>, IAsyncLifetime
    where T : DbContext
 {
    private SqlServerTestDbContextProvider<T>? _testCtxProvider;
@@ -16,8 +17,9 @@ public abstract class SqlServerDbContextIntegrationTests<T> : ITestDbContextProv
    /// </summary>
    protected SqlServerTestDbContextProvider<T> TestCtxProvider => _testCtxProvider ??= TestCtxProviderBuilder.Build();
 
-   private bool _isProviderConfigured;
    private readonly SqlServerTestDbContextProviderBuilder<T> _testCtxProviderBuilder;
+   private bool _isDisposed;
+   private bool _isProviderConfigured;
 
    /// <summary>
    /// Gets the <see cref="SqlServerTestDbContextProviderBuilder{T}"/> which is created on the first access.
@@ -97,23 +99,61 @@ public T CreateDbContext(bool useMasterConnection)
       return TestCtxProvider.CreateDbContext(useMasterConnection);
    }
 
+   /// <inheritdoc />
+   public virtual Task InitializeAsync()
+   {
+      return Task.CompletedTask;
+   }
+
    /// <inheritdoc />
    public void Dispose()
    {
+      if (_isDisposed)
+         return;
+
+      _isDisposed = true;
+
       Dispose(true);
       GC.SuppressFinalize(this);
    }
 
+   /// <inheritdoc />
+   async Task IAsyncLifetime.DisposeAsync()
+   {
+      await DisposeAsync();
+   }
+
+   /// <inheritdoc />
+   public async ValueTask DisposeAsync()
+   {
+      if (_isDisposed)
+         return;
+
+      _isDisposed = true;
+
+      await DisposeAsync(true);
+      GC.SuppressFinalize(this);
+   }
+
    /// <summary>
    /// Disposes of managed resources like the <see cref="SqlServerTestDbContextProvider{T}"/>.
    /// </summary>
    /// <param name="disposing">Indication that the method is being called by <see cref="Dispose()"/>.</param>
    protected virtual void Dispose(bool disposing)
+   {
+      DisposeAsync(disposing).AsTask().ConfigureAwait(false).GetAwaiter().GetResult();
+   }
+
+   /// <summary>
+   /// Disposes of managed resources like the <see cref="SqlServerTestDbContextProvider{T}"/>.
+   /// </summary>
+   /// <param name="disposing">Indication that the method is being called by <see cref="Dispose()"/>.</param>
+   protected virtual async ValueTask DisposeAsync(bool disposing)
    {
       if (!disposing)
          return;
 
-      _testCtxProvider?.Dispose();
+      await (_testCtxProvider?.DisposeAsync() ?? ValueTask.CompletedTask);
       _testCtxProvider = null;
    }
 }

src/Thinktecture.EntityFrameworkCore.SqlServer.Testing/EntityFrameworkCore/Testing/SqlServerTestDbContextProvider.cs

@@ -32,6 +32,7 @@ public class SqlServerTestDbContextProvider<T> : ITestDbContextProvider<T>
    private IDbContextTransaction? _tx;
    private bool _isAtLeastOneContextCreated;
    private readonly IsolationLevel _sharedTablesIsolationLevel;
+   private bool _isDisposed;
 
    /// <inheritdoc />
    public T ArrangeDbContext => _arrangeDbContext ??= CreateDbContext(true);
@@ -214,55 +215,84 @@ protected virtual void RunMigrations(T ctx)
 
    /// <summary>
    /// Rollbacks transaction if shared tables are used
-   /// otherwise the migrations are rolled back and all tables, functions, views and the newly generated schema are deleted.
+   /// otherwise performs cleanup according to provided <see cref="ITestIsolationOptions"/>.
    /// </summary>
    public void Dispose()
    {
+      if (_isDisposed)
+         return;
+
+      _isDisposed = true;
+
       Dispose(true);
       GC.SuppressFinalize(this);
    }
 
+   /// <summary>
+   /// Rollbacks transaction if shared tables are used
+   /// otherwise performs cleanup according to provided <see cref="ITestIsolationOptions"/>.
+   /// </summary>
+   public async ValueTask DisposeAsync()
+   {
+      if (_isDisposed)
+         return;
+
+      _isDisposed = true;
+
+      await DisposeAsync(true);
+      GC.SuppressFinalize(this);
+   }
+
    /// <summary>
    /// Disposes of inner resources.
    /// </summary>
    /// <param name="disposing">Indication whether this method is being called by the method <see cref="SqlServerTestDbContextProvider{T}.Dispose()"/>.</param>
    protected virtual void Dispose(bool disposing)
+   {
+      DisposeAsync(disposing).AsTask().ConfigureAwait(false).GetAwaiter().GetResult();
+   }
+
+   /// <summary>
+   /// Disposes of inner resources.
+   /// </summary>
+   /// <param name="disposing">Indication whether this method is being called by the method <see cref="SqlServerTestDbContextProvider{T}.Dispose()"/>.</param>
+   protected virtual async ValueTask DisposeAsync(bool disposing)
    {
       if (!disposing)
          return;
 
       if (_isAtLeastOneContextCreated)
       {
-         DisposeContextsAndRollbackMigrations();
+         await DisposeContextsAndRollbackMigrationsAsync(default);
          _isAtLeastOneContextCreated = false;
       }
 
       _masterConnection.Dispose();
       _testingLoggingOptions.Dispose();
    }
 
-   private void DisposeContextsAndRollbackMigrations()
+   private async Task DisposeContextsAndRollbackMigrationsAsync(CancellationToken cancellationToken)
    {
-      _tx?.Rollback();
-      _tx?.Dispose();
+      if (_tx is not null)
+      {
+         await _tx.RollbackAsync(cancellationToken);
+         await _tx.DisposeAsync();
+      }
 
       if (_isolationOptions.NeedsCleanup)
       {
          // Create a new ctx as a last resort to rollback migrations and clean up the database
-         using var ctx = _actDbContext ?? _arrangeDbContext ?? _assertDbContext ?? CreateDbContext(_masterDbContextOptions, new DbDefaultSchema(Schema));
+         await using var ctx = _actDbContext ?? _arrangeDbContext ?? _assertDbContext ?? CreateDbContext(_masterDbContextOptions, new DbDefaultSchema(Schema));
 
-         _isolationOptions.Cleanup(ctx, Schema);
+         await _isolationOptions.CleanupAsync(ctx, Schema, cancellationToken);
       }
 
-      _arrangeDbContext?.Dispose();
-      _actDbContext?.Dispose();
-      _assertDbContext?.Dispose();
+      await (_arrangeDbContext?.DisposeAsync() ?? ValueTask.CompletedTask);
+      await (_actDbContext?.DisposeAsync() ?? ValueTask.CompletedTask);
+      await (_assertDbContext?.DisposeAsync() ?? ValueTask.CompletedTask);
 
       _arrangeDbContext = null;
       _actDbContext = null;
       _assertDbContext = null;
    }
-
-
-
 }