From ffef7cda2cf34d9b6071e701b2d85d4314600878 Mon Sep 17 00:00:00 2001
From: Dan H <dan.r.hartley@gmail.com>
Date: Fri, 25 Mar 2022 15:13:40 +0000
Subject: [PATCH] Allowing more decoration methods to support open generic
 types.

---
 src/Scrutor/Decoration/DecorationFactory.cs   |  52 +++----
 .../OpenGenericDecorationStrategy.cs          | 136 ++++++++++--------
 .../OpenGenericDecorationTests.cs             |  33 +----
 3 files changed, 103 insertions(+), 118 deletions(-)

diff --git a/src/Scrutor/Decoration/DecorationFactory.cs b/src/Scrutor/Decoration/DecorationFactory.cs
index fbf1597..efc2469 100644
--- a/src/Scrutor/Decoration/DecorationFactory.cs
+++ b/src/Scrutor/Decoration/DecorationFactory.cs
@@ -1,26 +1,26 @@
-using System;
-
-namespace Scrutor.Decoration
-{
-    internal static class DecorationFactory
-    {
-        public static Decoration Create(Type serviceType, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory)
-        {
-            IDecorationStrategy strategy;
-
-            if (serviceType.IsOpenGeneric() && decoratorType is not null &&  decoratorType.IsOpenGeneric())
-            {
-                strategy = new OpenGenericDecorationStrategy(serviceType, decoratorType);
-            }
-            else
-            {
-                strategy = new ClosedTypeDecorationStrategy(serviceType, decoratorType, decoratorFactory);
-            }
-
-            return new Decoration(strategy);
-        }
-
-        public static Decoration Create<TService>(Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory)
-            => new(new ClosedTypeDecorationStrategy(typeof(TService), decoratorType, decoratorFactory)); 
-    }
-}
+using System;
+
+namespace Scrutor.Decoration
+{
+    internal static class DecorationFactory
+    {
+        public static Decoration Create(Type serviceType, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory)
+        {
+            IDecorationStrategy strategy;
+
+            if (serviceType.IsOpenGeneric())
+            {
+                strategy = new OpenGenericDecorationStrategy(serviceType, decoratorType, decoratorFactory);
+            }
+            else
+            {
+                strategy = new ClosedTypeDecorationStrategy(serviceType, decoratorType, decoratorFactory);
+            }
+
+            return new Decoration(strategy);
+        }
+
+        public static Decoration Create<TService>(Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory)
+            => new(new ClosedTypeDecorationStrategy(typeof(TService), decoratorType, decoratorFactory)); 
+    }
+}
diff --git a/src/Scrutor/Decoration/OpenGenericDecorationStrategy.cs b/src/Scrutor/Decoration/OpenGenericDecorationStrategy.cs
index b13d936..77c64ef 100644
--- a/src/Scrutor/Decoration/OpenGenericDecorationStrategy.cs
+++ b/src/Scrutor/Decoration/OpenGenericDecorationStrategy.cs
@@ -1,62 +1,74 @@
-using Microsoft.Extensions.DependencyInjection;
-using System;
-
-namespace Scrutor.Decoration
-{
-    internal class OpenGenericDecorationStrategy : IDecorationStrategy
-    {
-        private readonly Type _serviceType;
-        private readonly Type _decoratorType;
-
-        public OpenGenericDecorationStrategy(Type serviceType, Type decoratorType)
-        {
-            _serviceType = serviceType;
-            _decoratorType = decoratorType;
-        }
-
-        public Type ServiceType => _serviceType;
-
-        public bool CanDecorate(Type serviceType)
-        {
-            var canHandle = serviceType.IsGenericType
-                && (!serviceType.IsGenericTypeDefinition)
-                && _serviceType.GetGenericTypeDefinition() == serviceType.GetGenericTypeDefinition()
-                && HasCompatibleGenericArguments(serviceType);
-
-            return canHandle;
-        }
-
-        public Func<IServiceProvider, object> CreateDecorator(ServiceDescriptor descriptor)
-        {
-            var genericArguments = descriptor.ServiceType.GetGenericArguments();
-            var closedDecorator = _decoratorType.MakeGenericType(genericArguments);
-
-            return DecoratorInstanceFactory.Default(descriptor, closedDecorator);
-        }
-
-        private bool HasCompatibleGenericArguments(Type serviceType)
-        {
-            var canHandle = false;
-
-            if (_decoratorType is null)
-            {
-                canHandle = true;
-            }
-            else
-            {
-                var genericArguments = serviceType.GetGenericArguments();
-
-                try
-                {
-                    _ = _decoratorType.MakeGenericType(genericArguments);
-                    canHandle = true;
-                }
-                catch (ArgumentException)
-                {
-                }
-            }
-
-            return canHandle;
-        }
-    }
-}
+using Microsoft.Extensions.DependencyInjection;
+using System;
+
+namespace Scrutor.Decoration
+{
+    internal sealed class OpenGenericDecorationStrategy : IDecorationStrategy
+    {
+        private readonly Type _serviceType;
+        private readonly Type? _decoratorType;
+        private readonly Func<object, IServiceProvider, object>? _decoratorFactory;
+
+        public OpenGenericDecorationStrategy(Type serviceType, Type? decoratorType, Func<object, IServiceProvider, object>? decoratorFactory)
+        {
+            _serviceType = serviceType;
+            _decoratorType = decoratorType;
+            _decoratorFactory = decoratorFactory;
+        }
+
+        public Type ServiceType => _serviceType;
+
+        public bool CanDecorate(Type serviceType)
+        {
+            var canHandle = serviceType.IsGenericType
+                && (!serviceType.IsGenericTypeDefinition)
+                && _serviceType.GetGenericTypeDefinition() == serviceType.GetGenericTypeDefinition()
+                && HasCompatibleGenericArguments(serviceType);
+
+            return canHandle;
+        }
+
+        public Func<IServiceProvider, object> CreateDecorator(ServiceDescriptor descriptor)
+        {
+            if (_decoratorType is not null)
+            {
+                var genericArguments = descriptor.ServiceType.GetGenericArguments();
+                var closedDecorator = _decoratorType.MakeGenericType(genericArguments);
+
+                return DecoratorInstanceFactory.Default(descriptor, closedDecorator);
+            }
+
+            if (_decoratorFactory is not null)
+            {
+                return DecoratorInstanceFactory.Custom(descriptor, _decoratorFactory);
+            }
+
+            throw new InvalidOperationException($"Both serviceType and decoratorFactory can not be null.");
+        }
+
+        private bool HasCompatibleGenericArguments(Type serviceType)
+        {
+            var canHandle = false;
+
+            if (_decoratorType is null)
+            {
+                canHandle = true;
+            }
+            else
+            {
+                var genericArguments = serviceType.GetGenericArguments();
+
+                try
+                {
+                    _ = _decoratorType.MakeGenericType(genericArguments);
+                    canHandle = true;
+                }
+                catch (ArgumentException)
+                {
+                }
+            }
+
+            return canHandle;
+        }
+    }
+}
diff --git a/test/Scrutor.Tests/OpenGenericDecorationTests.cs b/test/Scrutor.Tests/OpenGenericDecorationTests.cs
index 51f6541..200cdd7 100644
--- a/test/Scrutor.Tests/OpenGenericDecorationTests.cs
+++ b/test/Scrutor.Tests/OpenGenericDecorationTests.cs
@@ -128,27 +128,6 @@ public void DecorationFunctionsDoSupportOpenGenericType()
             {
                 sc => sc.Decorate(typeof(QueryHandler<,>), typeof(LoggingQueryHandler<,>)),
                 sc => sc.TryDecorate(typeof(QueryHandler<,>), typeof(LoggingQueryHandler<,>)),
-            };
-
-            foreach (var decorationFunction in allDecorationFunctions)
-            {
-                var provider = ConfigureProvider(services =>
-                {
-                    services.AddSingleton<QueryHandler<MyQuery, MyResult>, MyQueryHandler>();
-                    decorationFunction(services);
-                });
-
-                var instance = provider.GetRequiredService<QueryHandler<MyQuery, MyResult>>();
-                var decorator = Assert.IsType<LoggingQueryHandler<MyQuery, MyResult>>(instance);
-                Assert.IsType<MyQueryHandler>(decorator.Inner);
-            }
-        }
-
-        [Fact]
-        public void DecorationFunctionsDoNotSupportOpenGenericType()
-        {
-            var allDecorationFunctions = new Action<IServiceCollection>[]
-            {
                 sc => sc.Decorate(typeof(QueryHandler<,>), (object obj, IServiceProvider sp) => new LoggingQueryHandler<MyQuery, MyResult>((IQueryHandler<MyQuery, MyResult>)obj)),
                 sc => sc.TryDecorate(typeof(QueryHandler<,>), (object obj, IServiceProvider sp) => new LoggingQueryHandler<MyQuery, MyResult>((IQueryHandler<MyQuery, MyResult>)obj)),
                 sc => sc.Decorate(typeof(QueryHandler<,>), (object obj) => new LoggingQueryHandler<MyQuery, MyResult>((IQueryHandler<MyQuery, MyResult>)obj)),
@@ -160,18 +139,12 @@ public void DecorationFunctionsDoNotSupportOpenGenericType()
                 var provider = ConfigureProvider(services =>
                 {
                     services.AddSingleton<QueryHandler<MyQuery, MyResult>, MyQueryHandler>();
-
-                    try
-                    {
-                        decorationFunction(services);
-                    }
-                    catch (MissingTypeRegistrationException)
-                    {
-                    }
+                    decorationFunction(services);
                 });
 
                 var instance = provider.GetRequiredService<QueryHandler<MyQuery, MyResult>>();
-                _ = Assert.IsType<MyQueryHandler>(instance);
+                var decorator = Assert.IsType<LoggingQueryHandler<MyQuery, MyResult>>(instance);
+                Assert.IsType<MyQueryHandler>(decorator.Inner);
             }
         }