From 09211fff4f8eafeff003aa6bd216a770daab77f8 Mon Sep 17 00:00:00 2001
From: danielreynolds1 <55194784+danielreynolds1@users.noreply.github.com>
Date: Thu, 7 Nov 2024 10:07:45 +0000
Subject: [PATCH] Added "AddDataLoader<TService, TImplementation>" overload
 that takes a factory (#7674)

---
 .../DataLoaderServiceCollectionExtensions.cs  | 13 +++++
 ...aLoaderServiceCollectionExtensionsTests.cs | 56 +++++++++++++++++++
 2 files changed, 69 insertions(+)
 create mode 100644 src/GreenDonut/test/Core.Tests/DependencyInjection/DataLoaderServiceCollectionExtensionsTests.cs

diff --git a/src/GreenDonut/src/Core/DependencyInjection/DataLoaderServiceCollectionExtensions.cs b/src/GreenDonut/src/Core/DependencyInjection/DataLoaderServiceCollectionExtensions.cs
index 3e6f24daa58..337cf21ddb7 100644
--- a/src/GreenDonut/src/Core/DependencyInjection/DataLoaderServiceCollectionExtensions.cs
+++ b/src/GreenDonut/src/Core/DependencyInjection/DataLoaderServiceCollectionExtensions.cs
@@ -41,6 +41,19 @@ public static IServiceCollection AddDataLoader<T>(
         return services;
     }
 
+    public static IServiceCollection AddDataLoader<TService, TImplementation>(
+        this IServiceCollection services,
+        Func<IServiceProvider, TImplementation> factory)
+        where TService : class, IDataLoader
+        where TImplementation : class, TService
+    {
+        services.TryAddDataLoaderCore();
+        services.AddSingleton(new DataLoaderRegistration(typeof(TService), typeof(TImplementation), sp => factory(sp)));
+        services.TryAddScoped<TImplementation>(sp => sp.GetDataLoader<TImplementation>());
+        services.TryAddScoped<TService>(sp => sp.GetDataLoader<TService>());
+        return services;
+    }
+
     public static IServiceCollection TryAddDataLoaderCore(
         this IServiceCollection services)
     {
diff --git a/src/GreenDonut/test/Core.Tests/DependencyInjection/DataLoaderServiceCollectionExtensionsTests.cs b/src/GreenDonut/test/Core.Tests/DependencyInjection/DataLoaderServiceCollectionExtensionsTests.cs
new file mode 100644
index 00000000000..42bc1672d34
--- /dev/null
+++ b/src/GreenDonut/test/Core.Tests/DependencyInjection/DataLoaderServiceCollectionExtensionsTests.cs
@@ -0,0 +1,56 @@
+using GreenDonut;
+using Xunit;
+using static GreenDonut.TestHelpers;
+
+namespace Microsoft.Extensions.DependencyInjection;
+
+public class DataLoaderServiceCollectionExtensionsTests
+{
+    [Fact]
+    public void ImplFactoryIsCalledWhenServiceIsResolved()
+    {
+        // arrange
+        var factoryCalled = false;
+        var fetch = CreateFetch<string, string>();
+        var services = new ServiceCollection()
+            .AddScoped<IBatchScheduler, ManualBatchScheduler>()
+            .AddDataLoader(sp =>
+            {
+                factoryCalled = true;
+                return new DataLoader<string, string>(fetch, sp.GetRequiredService<IBatchScheduler>());
+            });
+        var scope = services.BuildServiceProvider().CreateScope();
+
+        // act
+        var dataLoader = scope.ServiceProvider.GetRequiredService<DataLoader<string, string>>();
+
+        // assert
+        Assert.NotNull(dataLoader);
+        Assert.True(factoryCalled);
+    }
+
+    [Fact]
+    public void InterfaceImplFactoryIsCalledWhenServiceIsResolved()
+    {
+        // arrange
+        var factoryCalled = false;
+        var fetch = CreateFetch<string, string>();
+        var services = new ServiceCollection()
+            .AddScoped<IBatchScheduler, ManualBatchScheduler>()
+            .AddDataLoader<IDataLoader<string, string>, DataLoader<string, string>>(sp =>
+            {
+                factoryCalled = true;
+                return new DataLoader<string, string>(fetch, sp.GetRequiredService<IBatchScheduler>());
+            });
+        var scope = services.BuildServiceProvider().CreateScope();
+
+        // act
+        var dataLoader = scope.ServiceProvider.GetRequiredService<DataLoader<string, string>>();
+        var asInterface = scope.ServiceProvider.GetRequiredService<IDataLoader<string, string>>();
+
+        // assert
+        Assert.NotNull(dataLoader);
+        Assert.NotNull(asInterface);
+        Assert.True(factoryCalled);
+    }
+}