From beab42841e2988007889ddc27a2bde70d0f30b1e Mon Sep 17 00:00:00 2001
From: Bogdan <mynameisbogdan@users.noreply.github.com>
Date: Mon, 27 Apr 2026 21:07:37 +0300
Subject: [PATCH] Add async overloaded methods for MigrationTest

---
 .../Framework/DirectDataMapper.cs             | 40 +++++++++++++++++++
 .../Framework/MigrationTest.cs                |  6 +++
 2 files changed, 46 insertions(+)

diff --git a/src/NzbDrone.Core.Test/Framework/DirectDataMapper.cs b/src/NzbDrone.Core.Test/Framework/DirectDataMapper.cs
index f09fdbc49..3b7db72a5 100644
--- a/src/NzbDrone.Core.Test/Framework/DirectDataMapper.cs
+++ b/src/NzbDrone.Core.Test/Framework/DirectDataMapper.cs
@@ -2,6 +2,8 @@
 using System.Collections.Generic;
 using System.Data;
 using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
 using NzbDrone.Common.Serializer;
 using NzbDrone.Core.Datastore;
 
@@ -10,9 +12,13 @@ namespace NzbDrone.Core.Test.Framework
     public interface IDirectDataMapper
     {
         List<Dictionary<string, object>> Query(string sql);
+        Task<List<Dictionary<string, object>>> QueryAsync(string sql, CancellationToken cancellationToken = default);
         List<T> Query<T>(string sql)
             where T : new();
+        Task<List<T>> QueryAsync<T>(string sql, CancellationToken cancellationToken = default)
+            where T : new();
         T QueryScalar<T>(string sql);
+        Task<T> QueryScalarAsync<T>(string sql, CancellationToken cancellationToken = default);
     }
 
     public class DirectDataMapper : IDirectDataMapper
@@ -38,6 +44,18 @@ public DataTable GetDataTable(string sql)
             }
         }
 
+        public async Task<DataTable> GetDataTableAsync(string sql, CancellationToken cancellationToken = default)
+        {
+            await using var connection = await _database.OpenConnectionAsync(cancellationToken);
+            await using var cmd = connection.CreateCommand();
+
+            var dataTable = new DataTable();
+            cmd.CommandText = sql;
+            dataTable.Load(await cmd.ExecuteReaderAsync(cancellationToken));
+
+            return dataTable;
+        }
+
         public List<Dictionary<string, object>> Query(string sql)
         {
             var dataTable = GetDataTable(sql);
@@ -45,6 +63,13 @@ public List<Dictionary<string, object>> Query(string sql)
             return dataTable.Rows.Cast<DataRow>().Select(MapToDictionary).ToList();
         }
 
+        public async Task<List<Dictionary<string, object>>> QueryAsync(string sql, CancellationToken cancellationToken = default)
+        {
+            var dataTable = await GetDataTableAsync(sql, cancellationToken);
+
+            return dataTable.Rows.Cast<DataRow>().Select(MapToDictionary).ToList();
+        }
+
         public List<T> Query<T>(string sql)
             where T : new()
         {
@@ -53,6 +78,14 @@ public List<T> Query<T>(string sql)
             return dataTable.Rows.Cast<DataRow>().Select(MapToObject<T>).ToList();
         }
 
+        public async Task<List<T>> QueryAsync<T>(string sql, CancellationToken cancellationToken = default)
+            where T : new()
+        {
+            var dataTable = await GetDataTableAsync(sql, cancellationToken);
+
+            return dataTable.Rows.Cast<DataRow>().Select(MapToObject<T>).ToList();
+        }
+
         public T QueryScalar<T>(string sql)
         {
             var dataTable = GetDataTable(sql);
@@ -60,6 +93,13 @@ public T QueryScalar<T>(string sql)
             return dataTable.Rows.Cast<DataRow>().Select(d => MapValue(d, 0, typeof(T))).Cast<T>().FirstOrDefault();
         }
 
+        public async Task<T> QueryScalarAsync<T>(string sql, CancellationToken cancellationToken = default)
+        {
+            var dataTable = await GetDataTableAsync(sql, cancellationToken);
+
+            return dataTable.Rows.Cast<DataRow>().Select(d => MapValue(d, 0, typeof(T))).Cast<T>().FirstOrDefault();
+        }
+
         protected Dictionary<string, object> MapToDictionary(DataRow dataRow)
         {
             var item = new Dictionary<string, object>();
diff --git a/src/NzbDrone.Core.Test/Framework/MigrationTest.cs b/src/NzbDrone.Core.Test/Framework/MigrationTest.cs
index 9ac581a83..bc3be78f8 100644
--- a/src/NzbDrone.Core.Test/Framework/MigrationTest.cs
+++ b/src/NzbDrone.Core.Test/Framework/MigrationTest.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Data;
+using System.Threading.Tasks;
 using FluentMigrator;
 using Microsoft.Extensions.Logging;
 using NLog.Extensions.Logging;
@@ -31,6 +32,11 @@ protected virtual IDbConnection WithDapperMigrationTestDb(Action<TMigration> bef
             return WithMigrationAction(beforeMigration).OpenConnection();
         }
 
+        protected virtual async Task<IDbConnection> WithDapperMigrationTestDbAsync(Action<TMigration> beforeMigration = null)
+        {
+            return await WithMigrationAction(beforeMigration).OpenConnectionAsync();
+        }
+
         protected override void SetupLogging()
         {
             Mocker.SetConstant<ILoggerProvider>(Mocker.Resolve<NLogLoggerProvider>());