Adding open stream behaviors

jbogard · Jul 6, 2023 · dd15bfd · dd15bfd
1 parent 4ba26db
commit dd15bfd
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 15 deletions.
diff --git a/src/MediatR/MicrosoftExtensionsDI/MediatrServiceConfiguration.cs b/src/MediatR/MicrosoftExtensionsDI/MediatrServiceConfiguration.cs
@@ -131,11 +131,11 @@ public MediatRServiceConfiguration AddOpenBehavior(Type openBehaviorType, Servic
         }
 
         var implementedGenericInterfaces = openBehaviorType.GetInterfaces().Where(i => i.IsGenericType).Select(i => i.GetGenericTypeDefinition());
-        var implementedOpenBehaviorInterfaces = new HashSet<Type>(implementedGenericInterfaces.Where(i => i == typeof(IPipelineBehavior<,>) || i == typeof(IStreamPipelineBehavior<,>)));
+        var implementedOpenBehaviorInterfaces = new HashSet<Type>(implementedGenericInterfaces.Where(i => i == typeof(IPipelineBehavior<,>)));
 
         if (implementedOpenBehaviorInterfaces.Count == 0)
         {
-            throw new InvalidOperationException($"{openBehaviorType.Name} must implement {typeof(IPipelineBehavior<,>).FullName} or {typeof(IStreamPipelineBehavior<,>).FullName}");
+            throw new InvalidOperationException($"{openBehaviorType.Name} must implement {typeof(IPipelineBehavior<,>).FullName}");
         }
 
         foreach (var openBehaviorInterface in implementedOpenBehaviorInterfaces)
@@ -169,4 +169,35 @@ public MediatRServiceConfiguration AddStreamBehavior(Type serviceType, Type impl
 
         return this;
     }
+
+    /// <summary>
+    /// Registers an open stream behavior type against the <see cref="IStreamPipelineBehavior{TRequest,TResponse}"/> open generic interface type
+    /// </summary>
+    /// <param name="openBehaviorType">An open generic stream behavior type</param>
+    /// <param name="serviceLifetime">Optional service lifetime, defaults to <see cref="ServiceLifetime.Transient"/>.</param>
+    /// <returns>This</returns>
+    public MediatRServiceConfiguration AddOpenStreamBehavior(Type openBehaviorType, ServiceLifetime serviceLifetime = ServiceLifetime.Transient)
+    {
+        if (!openBehaviorType.IsGenericType)
+        {
+            throw new InvalidOperationException($"{openBehaviorType.Name} must be generic");
+        }
+
+        var implementedGenericInterfaces = openBehaviorType.GetInterfaces().Where(i => i.IsGenericType).Select(i => i.GetGenericTypeDefinition());
+        var implementedOpenBehaviorInterfaces = new HashSet<Type>(implementedGenericInterfaces.Where(i => i == typeof(IStreamPipelineBehavior<,>)));
+
+        if (implementedOpenBehaviorInterfaces.Count == 0)
+        {
+            throw new InvalidOperationException($"{openBehaviorType.Name} must implement {typeof(IStreamPipelineBehavior<,>).FullName}");
+        }
+
+        foreach (var openBehaviorInterface in implementedOpenBehaviorInterfaces)
+        {
+            StreamBehaviorsToRegister.Add(new ServiceDescriptor(openBehaviorInterface, openBehaviorType, serviceLifetime));
+        }
+
+        return this;
+    }
+
+
 }
diff --git a/test/MediatR.Tests/MicrosoftExtensionsDI/PipelineTests.cs b/test/MediatR.Tests/MicrosoftExtensionsDI/PipelineTests.cs
@@ -582,41 +582,44 @@ public void Should_handle_open_behavior_registration()
     {
         var cfg = new MediatRServiceConfiguration();
         cfg.AddOpenBehavior(typeof(OpenBehavior<,>));
-        cfg.AddOpenBehavior(typeof(OpenStreamBehavior<,>));
+        cfg.AddOpenStreamBehavior(typeof(OpenStreamBehavior<,>));
 
-        cfg.BehaviorsToRegister.Count.ShouldBe(2);
+        cfg.BehaviorsToRegister.Count.ShouldBe(1);
+        cfg.StreamBehaviorsToRegister.Count.ShouldBe(1);
 
         cfg.BehaviorsToRegister[0].ServiceType.ShouldBe(typeof(IPipelineBehavior<,>));
         cfg.BehaviorsToRegister[0].ImplementationType.ShouldBe(typeof(OpenBehavior<,>));
         cfg.BehaviorsToRegister[0].ImplementationFactory.ShouldBeNull();
         cfg.BehaviorsToRegister[0].ImplementationInstance.ShouldBeNull();
         cfg.BehaviorsToRegister[0].Lifetime.ShouldBe(ServiceLifetime.Transient);
 
-        cfg.BehaviorsToRegister[1].ServiceType.ShouldBe(typeof(IStreamPipelineBehavior<,>));
-        cfg.BehaviorsToRegister[1].ImplementationType.ShouldBe(typeof(OpenStreamBehavior<,>));
-        cfg.BehaviorsToRegister[1].ImplementationFactory.ShouldBeNull();
-        cfg.BehaviorsToRegister[1].ImplementationInstance.ShouldBeNull();
-        cfg.BehaviorsToRegister[1].Lifetime.ShouldBe(ServiceLifetime.Transient);
+        cfg.StreamBehaviorsToRegister[0].ServiceType.ShouldBe(typeof(IStreamPipelineBehavior<,>));
+        cfg.StreamBehaviorsToRegister[0].ImplementationType.ShouldBe(typeof(OpenStreamBehavior<,>));
+        cfg.StreamBehaviorsToRegister[0].ImplementationFactory.ShouldBeNull();
+        cfg.StreamBehaviorsToRegister[0].ImplementationInstance.ShouldBeNull();
+        cfg.StreamBehaviorsToRegister[0].Lifetime.ShouldBe(ServiceLifetime.Transient);
     }
 
     [Fact]
     public void Should_handle_open_behaviors_registration_from_a_single_type()
     {
         var cfg = new MediatRServiceConfiguration();
         cfg.AddOpenBehavior(typeof(MultiOpenBehavior<,>), ServiceLifetime.Singleton);
+        cfg.AddOpenStreamBehavior(typeof(MultiOpenBehavior<,>), ServiceLifetime.Singleton);
 
-        cfg.BehaviorsToRegister.Count.ShouldBe(2);
+        cfg.BehaviorsToRegister.Count.ShouldBe(1);
+        cfg.StreamBehaviorsToRegister.Count.ShouldBe(1);
 
         cfg.BehaviorsToRegister[0].ServiceType.ShouldBe(typeof(IPipelineBehavior<,>));
         cfg.BehaviorsToRegister[0].ImplementationType.ShouldBe(typeof(MultiOpenBehavior<,>));
         cfg.BehaviorsToRegister[0].ImplementationFactory.ShouldBeNull();
         cfg.BehaviorsToRegister[0].ImplementationInstance.ShouldBeNull();
         cfg.BehaviorsToRegister[0].Lifetime.ShouldBe(ServiceLifetime.Singleton);
 
-        cfg.BehaviorsToRegister[1].ServiceType.ShouldBe(typeof(IStreamPipelineBehavior<,>));
-        cfg.BehaviorsToRegister[1].ImplementationType.ShouldBe(typeof(MultiOpenBehavior<,>));
-        cfg.BehaviorsToRegister[1].ImplementationFactory.ShouldBeNull();
-        cfg.BehaviorsToRegister[1].ImplementationInstance.ShouldBeNull();
-        cfg.BehaviorsToRegister[1].Lifetime.ShouldBe(ServiceLifetime.Singleton);
+        cfg.StreamBehaviorsToRegister[0].ServiceType.ShouldBe(typeof(IStreamPipelineBehavior<,>));
+        cfg.StreamBehaviorsToRegister[0].ImplementationType.ShouldBe(typeof(MultiOpenBehavior<,>));
+        cfg.StreamBehaviorsToRegister[0].ImplementationFactory.ShouldBeNull();
+        cfg.StreamBehaviorsToRegister[0].ImplementationInstance.ShouldBeNull();
+        cfg.StreamBehaviorsToRegister[0].Lifetime.ShouldBe(ServiceLifetime.Singleton);
     }
 }