ValueTask a bunch of Database Functions

src/Tgstation.Server.Host/Components/Instance.cs

@@ -474,7 +474,8 @@ namespace Tgstation.Server.Host.Components
 							await repo.ResetToSha(startSha, progressReporter, CancellationToken.None);
 							throw;
 						}
-				});
+				})
+			.AsTask();
 #pragma warning restore CA1502   // Cyclomatic complexity
 
 		/// <summary>

src/Tgstation.Server.Host/Components/InstanceManager.cs

@@ -258,7 +258,7 @@ namespace Tgstation.Server.Host.Components
 					logger.LogDebug("Reverting instance {instanceId}'s path to {oldPath} in the DB...", instance.Id, oldPath);
 
 					// DCT: Operation must always run
-					await databaseContextFactory.UseContext(db =>
+					await databaseContextFactory.UseContext2(db =>
 					{
 						var targetInstance = new Models.Instance
 						{
@@ -544,8 +544,9 @@ namespace Tgstation.Server.Host.Components
 				await InitializeSwarm(cancellationToken);
 
 				List<Models.Instance> dbInstances = null;
-				var instanceEnumeration = databaseContextFactory.UseContext(
-					async databaseContext => dbInstances = await databaseContext
+
+				async ValueTask EnumerateInstances(IDatabaseContext databaseContext)
+					=> dbInstances = await databaseContext
 						.Instances
 						.AsQueryable()
 						.Where(x => x.Online.Value && x.SwarmIdentifer == swarmConfiguration.Identifier)
@@ -553,12 +554,14 @@ namespace Tgstation.Server.Host.Components
 						.Include(x => x.ChatSettings)
 							.ThenInclude(x => x.Channels)
 						.Include(x => x.DreamDaemonSettings)
-						.ToListAsync(cancellationToken));
+						.ToListAsync(cancellationToken);
+
+				var instanceEnumeration = databaseContextFactory.UseContext(EnumerateInstances);
 
 				var factoryStartup = instanceFactory.StartAsync(cancellationToken);
 				var jobManagerStartup = jobService.StartAsync(cancellationToken);
 
-				await Task.WhenAll(instanceEnumeration, factoryStartup, jobManagerStartup);
+				await Task.WhenAll(instanceEnumeration.AsTask(), factoryStartup, jobManagerStartup);
 
 				var instanceOnliningTasks = dbInstances.Select(
 					async metadata =>

src/Tgstation.Server.Host/Components/Repository/RepositoryUpdateService.cs

@@ -192,7 +192,7 @@ namespace Tgstation.Server.Host.Components.Repository
 				Id = instanceId,
 			};
 
-			Task CallLoadRevInfo(Models.TestMerge testMergeToAdd = null, string lastOriginCommitSha = null) => databaseContextFactory
+			ValueTask CallLoadRevInfo(Models.TestMerge testMergeToAdd = null, string lastOriginCommitSha = null) => databaseContextFactory
 				.UseContext(
 					async databaseContext =>
 					{
@@ -243,7 +243,7 @@ namespace Tgstation.Server.Host.Components.Repository
 			await CallLoadRevInfo();
 
 			// apply new rev info, tracking applied test merges
-			Task UpdateRevInfo(Models.TestMerge testMergeToAdd = null) => CallLoadRevInfo(testMergeToAdd, lastRevisionInfo.OriginCommitSha);
+			ValueTask UpdateRevInfo(Models.TestMerge testMergeToAdd = null) => CallLoadRevInfo(testMergeToAdd, lastRevisionInfo.OriginCommitSha);
 
 			try
 			{

src/Tgstation.Server.Host/Components/Session/ISessionPersistor.cs

@@ -13,8 +13,8 @@ namespace Tgstation.Server.Host.Components.Session
 		/// </summary>
 		/// <param name="reattachInformation">The <see cref="ReattachInformation"/> to save.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		Task Save(ReattachInformation reattachInformation, CancellationToken cancellationToken);
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		ValueTask Save(ReattachInformation reattachInformation, CancellationToken cancellationToken);
 
 		/// <summary>
 		/// Load a saved <see cref="ReattachInformation"/>.
@@ -27,7 +27,7 @@ namespace Tgstation.Server.Host.Components.Session
 		/// Clear any stored <see cref="ReattachInformation"/>.
 		/// </summary>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		Task Clear(CancellationToken cancellationToken);
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		ValueTask Clear(CancellationToken cancellationToken);
 	}
 }

src/Tgstation.Server.Host/Components/Session/SessionPersistor.cs

@@ -65,7 +65,7 @@ namespace Tgstation.Server.Host.Components.Session
 		}
 
 		/// <inheritdoc />
-		public Task Save(ReattachInformation reattachInformation, CancellationToken cancellationToken) => databaseContextFactory.UseContext(async (db) =>
+		public ValueTask Save(ReattachInformation reattachInformation, CancellationToken cancellationToken) => databaseContextFactory.UseContext(async (db) =>
 		{
 			ArgumentNullException.ThrowIfNull(reattachInformation);
 
@@ -212,7 +212,7 @@ namespace Tgstation.Server.Host.Components.Session
 		}
 
 		/// <inheritdoc />
-		public Task Clear(CancellationToken cancellationToken) => databaseContextFactory
+		public ValueTask Clear(CancellationToken cancellationToken) => databaseContextFactory
 			.UseContext(
 				db =>
 				{
@@ -226,8 +226,8 @@ namespace Tgstation.Server.Host.Components.Session
 		/// <param name="databaseContext">The <see cref="IDatabaseContext"/> to use.</param>
 		/// <param name="instant">If an SQL DELETE WHERE command should be used rather than an Entity Framework transaction.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		async Task ClearImpl(IDatabaseContext databaseContext, bool instant, CancellationToken cancellationToken)
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		async ValueTask ClearImpl(IDatabaseContext databaseContext, bool instant, CancellationToken cancellationToken)
 		{
 			var baseQuery = databaseContext
 				.ReattachInformations

src/Tgstation.Server.Host/Components/Watchdog/BasicWatchdog.cs

@@ -270,11 +270,9 @@ namespace Tgstation.Server.Host.Components.Watchdog
 		/// Called to save the current <see cref="Server"/> into the <see cref="WatchdogBase.SessionPersistor"/> when initially launched.
 		/// </summary>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		protected virtual Task SessionStartupPersist(CancellationToken cancellationToken)
-		{
-			return SessionPersistor.Save(Server.ReattachInformation, cancellationToken);
-		}
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		protected virtual ValueTask SessionStartupPersist(CancellationToken cancellationToken)
+			=> SessionPersistor.Save(Server.ReattachInformation, cancellationToken);
 
 		/// <summary>
 		/// Handler for <see cref="MonitorActivationReason.ActiveServerRebooted"/> when the <see cref="RebootState"/> is <see cref="RebootState.Normal"/>.

src/Tgstation.Server.Host/Components/Watchdog/WindowsWatchdog.cs

@@ -244,7 +244,7 @@ namespace Tgstation.Server.Host.Components.Watchdog
 		}
 
 		/// <inheritdoc />
-		protected override async Task SessionStartupPersist(CancellationToken cancellationToken)
+		protected override async ValueTask SessionStartupPersist(CancellationToken cancellationToken)
 		{
 			await ApplyInitialDmb(cancellationToken);
 			await base.SessionStartupPersist(cancellationToken);

src/Tgstation.Server.Host/Database/DatabaseContext.cs

@@ -283,7 +283,7 @@ namespace Tgstation.Server.Host.Database
 		public Task Drop(CancellationToken cancellationToken) => Database.EnsureDeletedAsync(cancellationToken);
 
 		/// <inheritdoc />
-		public async Task<bool> Migrate(ILogger<DatabaseContext> logger, CancellationToken cancellationToken)
+		public async ValueTask<bool> Migrate(ILogger<DatabaseContext> logger, CancellationToken cancellationToken)
 		{
 			ArgumentNullException.ThrowIfNull(logger);
 			var migrations = await Database.GetAppliedMigrationsAsync(cancellationToken);
@@ -396,7 +396,7 @@ namespace Tgstation.Server.Host.Database
 
 		/// <inheritdoc />
 #pragma warning disable CA1502 // Cyclomatic complexity
-		public async Task SchemaDowngradeForServerVersion(
+		public async ValueTask SchemaDowngradeForServerVersion(
 			ILogger<DatabaseContext> logger,
 			Version targetVersion,
 			DatabaseType currentDatabaseType,

src/Tgstation.Server.Host/Database/DatabaseContextFactory.cs

@@ -26,7 +26,16 @@ namespace Tgstation.Server.Host.Database
 		}
 
 		/// <inheritdoc />
-		public async Task UseContext(Func<IDatabaseContext, Task> operation)
+		public async ValueTask UseContext(Func<IDatabaseContext, ValueTask> operation)
+		{
+			ArgumentNullException.ThrowIfNull(operation);
+
+			await using var scope = scopeFactory.CreateAsyncScope();
+			await operation(scope.ServiceProvider.GetRequiredService<IDatabaseContext>());
+		}
+
+		/// <inheritdoc />
+		public async ValueTask UseContext2(Func<IDatabaseContext, Task> operation)
 		{
 			ArgumentNullException.ThrowIfNull(operation);
 

src/Tgstation.Server.Host/Database/DatabaseSeeder.cs

@@ -102,7 +102,7 @@ namespace Tgstation.Server.Host.Database
 		}
 
 		/// <inheritdoc />
-		public async Task Initialize(IDatabaseContext databaseContext, CancellationToken cancellationToken)
+		public async ValueTask Initialize(IDatabaseContext databaseContext, CancellationToken cancellationToken)
 		{
 			ArgumentNullException.ThrowIfNull(databaseContext);
 
@@ -131,7 +131,7 @@ namespace Tgstation.Server.Host.Database
 		}
 
 		/// <inheritdoc />
-		public Task Downgrade(IDatabaseContext databaseContext, Version downgradeVersion, CancellationToken cancellationToken)
+		public ValueTask Downgrade(IDatabaseContext databaseContext, Version downgradeVersion, CancellationToken cancellationToken)
 		{
 			ArgumentNullException.ThrowIfNull(databaseContext);
 			ArgumentNullException.ThrowIfNull(downgradeVersion);
@@ -172,8 +172,8 @@ namespace Tgstation.Server.Host.Database
 		/// </summary>
 		/// <param name="databaseContext">The <see cref="IDatabaseContext"/> to seed.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		async Task SeedDatabase(IDatabaseContext databaseContext, CancellationToken cancellationToken)
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		async ValueTask SeedDatabase(IDatabaseContext databaseContext, CancellationToken cancellationToken)
 		{
 			var adminUser = SeedAdminUser(databaseContext);
 
@@ -191,8 +191,8 @@ namespace Tgstation.Server.Host.Database
 		/// </summary>
 		/// <param name="databaseContext">The <see cref="IDatabaseContext"/> to sanitize.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		async Task SanitizeDatabase(IDatabaseContext databaseContext, CancellationToken cancellationToken)
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		async ValueTask SanitizeDatabase(IDatabaseContext databaseContext, CancellationToken cancellationToken)
 		{
 			var admin = await GetAdminUser(databaseContext, cancellationToken);
 			if (admin != null)
@@ -270,8 +270,8 @@ namespace Tgstation.Server.Host.Database
 		/// </summary>
 		/// <param name="databaseContext">The <see cref="IDatabaseContext"/> to reset the admin password for.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		async Task ResetAdminPassword(IDatabaseContext databaseContext, CancellationToken cancellationToken)
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		async ValueTask ResetAdminPassword(IDatabaseContext databaseContext, CancellationToken cancellationToken)
 		{
 			var admin = await GetAdminUser(databaseContext, cancellationToken);
 			if (admin != null)
@@ -302,8 +302,8 @@ namespace Tgstation.Server.Host.Database
 		/// </summary>
 		/// <param name="databaseContext">The <see cref="IDatabaseContext"/> to use.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task{TResult}"/> resulting in the admin <see cref="User"/> or <see langword="null"/>. If <see langword="null"/>, <see cref="IDatabaseContext.Save(CancellationToken)"/> must be called on <paramref name="databaseContext"/>.</returns>
-		async Task<User> GetAdminUser(IDatabaseContext databaseContext, CancellationToken cancellationToken)
+		/// <returns>A <see cref="ValueTask{TResult}"/> resulting in the admin <see cref="User"/> or <see langword="null"/>. If <see langword="null"/>, <see cref="IDatabaseContext.Save(CancellationToken)"/> must be called on <paramref name="databaseContext"/>.</returns>
+		async ValueTask<User> GetAdminUser(IDatabaseContext databaseContext, CancellationToken cancellationToken)
 		{
 			var admin = await databaseContext
 				.Users

src/Tgstation.Server.Host/Database/IDatabaseContext.cs

@@ -109,8 +109,8 @@ namespace Tgstation.Server.Host.Database
 		/// </summary>
 		/// <param name="logger">The <see cref="DatabaseContext"/> <see cref="ILogger"/> to use.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task{TResult}"/> resulting in <see langword="true"/> if the database should be seeded, <see langword="false"/> otherwise.</returns>
-		Task<bool> Migrate(ILogger<DatabaseContext> logger, CancellationToken cancellationToken);
+		/// <returns>A <see cref="ValueTask{TResult}"/> resulting in <see langword="true"/> if the database should be seeded, <see langword="false"/> otherwise.</returns>
+		ValueTask<bool> Migrate(ILogger<DatabaseContext> logger, CancellationToken cancellationToken);
 
 		/// <summary>
 		/// Attempt to downgrade the schema to the migration used for a given server <paramref name="targetVersion"/>.
@@ -119,8 +119,8 @@ namespace Tgstation.Server.Host.Database
 		/// <param name="targetVersion">The tgstation-server <see cref="Version"/> that the schema should downgrade for.</param>
 		/// <param name="currentDatabaseType">The <see cref="DatabaseType"/> in use.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		Task SchemaDowngradeForServerVersion(
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		ValueTask SchemaDowngradeForServerVersion(
 			ILogger<DatabaseContext> logger,
 			Version targetVersion,
 			DatabaseType currentDatabaseType,

src/Tgstation.Server.Host/Database/IDatabaseContextFactory.cs

@@ -12,7 +12,14 @@ namespace Tgstation.Server.Host.Database
 		/// Run an <paramref name="operation"/> in the scope of an <see cref="IDatabaseContext"/>.
 		/// </summary>
 		/// <param name="operation">The operation to run.</param>
-		/// <returns>A <see cref="Task"/> representing the running <paramref name="operation"/>.</returns>
-		Task UseContext(Func<IDatabaseContext, Task> operation);
+		/// <returns>A <see cref="ValueTask"/> representing the running <paramref name="operation"/>.</returns>
+		ValueTask UseContext(Func<IDatabaseContext, ValueTask> operation);
+
+		/// <summary>
+		/// Run an <paramref name="operation"/> in the scope of an <see cref="IDatabaseContext"/>.
+		/// </summary>
+		/// <param name="operation">The operation to run.</param>
+		/// <returns>A <see cref="ValueTask"/> representing the running <paramref name="operation"/>.</returns>
+		ValueTask UseContext2(Func<IDatabaseContext, Task> operation);
 	}
 }

src/Tgstation.Server.Host/Database/IDatabaseSeeder.cs

@@ -14,8 +14,8 @@ namespace Tgstation.Server.Host.Database
 		/// </summary>
 		/// <param name="databaseContext">The <see cref="IDatabaseContext"/> to setup.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		Task Initialize(IDatabaseContext databaseContext, CancellationToken cancellationToken);
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		ValueTask Initialize(IDatabaseContext databaseContext, CancellationToken cancellationToken);
 
 		/// <summary>
 		/// Migrate a given <paramref name="databaseContext"/> down.
@@ -23,7 +23,7 @@ namespace Tgstation.Server.Host.Database
 		/// <param name="databaseContext">The <see cref="IDatabaseContext"/> to downgrade.</param>
 		/// <param name="downgradeVersion">The migration <see cref="Version"/> to downgrade the <paramref name="databaseContext"/> to.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing the running operation.</returns>
-		Task Downgrade(IDatabaseContext databaseContext, Version downgradeVersion, CancellationToken cancellationToken);
+		/// <returns>A <see cref="ValueTask"/> representing the running operation.</returns>
+		ValueTask Downgrade(IDatabaseContext databaseContext, Version downgradeVersion, CancellationToken cancellationToken);
 	}
 }

src/Tgstation.Server.Host/Jobs/IJobManager.cs

@@ -23,8 +23,8 @@ namespace Tgstation.Server.Host.Jobs
 		/// <param name="job">The <see cref="Job"/>. Should at least have <see cref="Job.Instance"/> and <see cref="Api.Models.Internal.Job.Description"/>. If <see cref="Job.StartedBy"/> is <see langword="null"/>, the TGS user will be used.</param>
 		/// <param name="operation">The <see cref="JobEntrypoint"/> for the <paramref name="job"/>.</param>
 		/// <param name="cancellationToken">The <see cref="CancellationToken"/> for the operation.</param>
-		/// <returns>A <see cref="Task"/> representing a running operation.</returns>
-		Task RegisterOperation(Job job, JobEntrypoint operation, CancellationToken cancellationToken);
+		/// <returns>A <see cref="ValueTask"/> representing a running operation.</returns>
+		ValueTask RegisterOperation(Job job, JobEntrypoint operation, CancellationToken cancellationToken);
 
 		/// <summary>
 		/// Wait for a given <paramref name="job"/> to complete.

src/Tgstation.Server.Host/Jobs/JobService.cs

@@ -87,7 +87,7 @@ namespace Tgstation.Server.Host.Jobs
 		}
 
 		/// <inheritdoc />
-		public Task RegisterOperation(Job job, JobEntrypoint operation, CancellationToken cancellationToken)
+		public ValueTask RegisterOperation(Job job, JobEntrypoint operation, CancellationToken cancellationToken)
 			=> databaseContextFactory.UseContext(
 				async databaseContext =>
 				{
@@ -168,7 +168,8 @@ namespace Tgstation.Server.Host.Jobs
 				}
 
 				noMoreJobsShouldStart = false;
-			});
+			})
+			.AsTask();
 
 		/// <inheritdoc />
 		public async Task StopAsync(CancellationToken cancellationToken)

tests/Tgstation.Server.Host.Tests/Components/Chat/Providers/TestDiscordProvider.cs

@@ -36,7 +36,7 @@ namespace Tgstation.Server.Host.Components.Chat.Providers.Tests
 			mockSetup
 				.Setup(x => x.RegisterOperation(It.IsNotNull<Job>(), It.IsNotNull<JobEntrypoint>(), It.IsAny<CancellationToken>()))
 				.Callback<Job, JobEntrypoint, CancellationToken>((job, entrypoint, cancellationToken) => job.StartedBy ??= new User { })
-				.Returns(Task.CompletedTask);
+				.Returns(ValueTask.CompletedTask);
 			mockSetup
 				.Setup(x => x.WaitForJobCompletion(It.IsNotNull<Job>(), It.IsAny<User>(), It.IsAny<CancellationToken>(), It.IsAny<CancellationToken>()))
 				.Returns(Task.CompletedTask);

tests/Tgstation.Server.Host.Tests/Components/Chat/Providers/TestIrcProvider.cs

@@ -73,7 +73,7 @@ namespace Tgstation.Server.Host.Components.Chat.Providers.Tests
 			mockSetup
 				.Setup(x => x.RegisterOperation(It.IsNotNull<Job>(), It.IsNotNull<JobEntrypoint>(), It.IsAny<CancellationToken>()))
 				.Callback<Job, JobEntrypoint, CancellationToken>((job, entrypoint, cancellationToken) => job.StartedBy ??= new User { })
-				.Returns(Task.CompletedTask);
+				.Returns(ValueTask.CompletedTask);
 			mockSetup
 				.Setup(x => x.WaitForJobCompletion(It.IsNotNull<Job>(), It.IsAny<User>(), It.IsAny<CancellationToken>(), It.IsAny<CancellationToken>()))
 				.Returns(Task.CompletedTask);

tests/Tgstation.Server.Host.Tests/Database/TestDatabaseContextFactory.cs

@@ -1,4 +1,4 @@
-using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using Moq;
 using System;
@@ -41,12 +41,12 @@ namespace Tgstation.Server.Host.Database.Tests
 
 			var factory = new DatabaseContextFactory(mockScopeFactory.Object);
 
-			await Assert.ThrowsExceptionAsync<ArgumentNullException>(() => factory.UseContext(null));
+			await Assert.ThrowsExceptionAsync<ArgumentNullException>(async () => await factory.UseContext(null));
 
 			await factory.UseContext(context =>
 			{
 				Assert.AreSame(mockDbo, context);
-				return Task.CompletedTask;
+				return ValueTask.CompletedTask;
 			});
 
 			mockScopeFactory.VerifyAll();

tests/Tgstation.Server.Host.Tests/Swarm/TestableSwarmNode.cs

@@ -51,7 +51,7 @@ namespace Tgstation.Server.Host.Swarm.Tests
 		readonly Mock<IHttpClient> mockHttpClient;
 		readonly Mock<IDatabaseContextFactory> mockDBContextFactory;
 		readonly Mock<IDatabaseSeeder> mockDatabaseSeeder;
-		readonly ISetup<IDatabaseSeeder, Task> mockDatabaseSeederInitialize;
+		readonly ISetup<IDatabaseSeeder, ValueTask> mockDatabaseSeederInitialize;
 
 		readonly Action recreateControllerAndService;
 
@@ -95,7 +95,7 @@ namespace Tgstation.Server.Host.Swarm.Tests
 			mockDatabaseSeederInitialize = new Mock<IDatabaseSeeder>().Setup(x => x.Initialize(mockDatabaseContext, It.IsAny<CancellationToken>()));
 			mockDBContextFactory = new Mock<IDatabaseContextFactory>();
 			mockDBContextFactory
-				.Setup(x => x.UseContext(It.IsNotNull<Func<IDatabaseContext, Task>>()))
+				.Setup(x => x.UseContext(It.IsNotNull<Func<IDatabaseContext, ValueTask>>()))
 				.Callback<Func<IDatabaseContext, Task>>((func) => func(mockDatabaseContext));
 
 			var mockHttpClientFactory = new Mock<IAbstractHttpClientFactory>();
@@ -226,9 +226,13 @@ namespace Tgstation.Server.Host.Swarm.Tests
 				Assert.Fail("Initialized twice!");
 
 			if (!cancel)
-				mockDatabaseSeederInitialize.Returns(Task.CompletedTask).Verifiable();
+				mockDatabaseSeederInitialize.Returns(ValueTask.CompletedTask).Verifiable();
 			else
-				mockDatabaseSeederInitialize.ThrowsAsync(new TaskCanceledException()).Verifiable();
+				mockDatabaseSeederInitialize.Returns(async () =>
+				{
+					await Task.Yield();
+					throw new TaskCanceledException();
+				}).Verifiable();
 
 			Task<SwarmRegistrationResult> Invoke() => Service.Initialize(default);
 

tests/Tgstation.Server.Tests/TestDatabase.cs

@@ -73,7 +73,7 @@ namespace Tgstation.Server.Tests
 				return null;
 			}
 
-			using var context = CreateContext();
+			await using var context = CreateContext();
 			await context.Database.EnsureDeletedAsync();
 			await context.Database.MigrateAsync(default);