From 34068b484dac0e5fabe45899e108a6d0412c44c7 Mon Sep 17 00:00:00 2001
From: "Jeremy D. Miller" <jeremydmiller@yahoo.com>
Date: Wed, 18 Dec 2024 14:27:50 -0600
Subject: [PATCH] Fixed the codegen ahead of time with sticky handlers. Closes
 GH-1163

---
 .../Acceptance/sticky_message_handlers.cs     | 24 +++++++++++++++++++
 src/Testing/CoreTests/IntegrationContext.cs   |  9 ++++++-
 .../Runtime/Handlers/HandlerChain.cs          | 20 +++++++++++-----
 .../Handlers/HandlerGraph.GeneratesCode.cs    | 21 +++++++++++++++-
 4 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/src/Testing/CoreTests/Acceptance/sticky_message_handlers.cs b/src/Testing/CoreTests/Acceptance/sticky_message_handlers.cs
index 0c00d829d..623db386a 100644
--- a/src/Testing/CoreTests/Acceptance/sticky_message_handlers.cs
+++ b/src/Testing/CoreTests/Acceptance/sticky_message_handlers.cs
@@ -1,7 +1,9 @@
 using System.Collections;
 using JasperFx.CodeGeneration;
+using JasperFx.CodeGeneration.Model;
 using JasperFx.Core;
 using JasperFx.Core.Reflection;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Wolverine.ComplianceTests;
@@ -35,6 +37,28 @@ public async Task send_message_is_handled_by_both_handlers_independently_by_attr
         records.ShouldContain(new StickyMessageResponse("blue", stickyMessage, new Uri("local://blue")));
     }
 
+    [Fact]
+    public void generate_code_with_sticky_handlers()
+    {
+        /*
+         * Steps -->
+         * HandlerChain needs to have some kind of suffix added to the generated handler type to mark by the endpoint. Use the full handler type name maybe
+         * HandlerChain if sticky needs to disambiguate the file name too
+         * HandlerGraph implementation of the getting files needs to change
+         * 
+         */
+        
+        
+        var collections = Host.Services.GetServices<ICodeFileCollection>().ToArray();
+
+        var builder = new DynamicCodeBuilder(Host.Services, collections)
+        {
+            ServiceVariableSource = Host.Services.GetService<IServiceVariableSource>()
+        };
+
+        builder.GenerateAllCode();
+    }
+
     public class FakeChainPolicy : IChainPolicy
     {
         public List<HandlerChain> Chains = new();
diff --git a/src/Testing/CoreTests/IntegrationContext.cs b/src/Testing/CoreTests/IntegrationContext.cs
index f3bf167f0..50a2da742 100644
--- a/src/Testing/CoreTests/IntegrationContext.cs
+++ b/src/Testing/CoreTests/IntegrationContext.cs
@@ -1,7 +1,11 @@
-using JasperFx.Core.Reflection;
+using CoreTests.Acceptance;
+using CoreTests.Bugs;
+using JasperFx.Core.Reflection;
 using Lamar.Microsoft.DependencyInjection;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
+using NSubstitute;
 using Oakton.Resources;
 using Wolverine.ComplianceTests;
 using Wolverine.ComplianceTests.Compliance;
@@ -25,6 +29,9 @@ public DefaultApp()
                 
                 opts.IncludeType<MessageConsumer>();
                 opts.IncludeType<InvokedMessageHandler>();
+
+                opts.Services.AddSingleton(Substitute.For<IIdentityService>());
+                opts.Services.AddSingleton(Substitute.For<ITrackedTaskRepository>());
             })
             .UseResourceSetupOnStartup(StartupAction.ResetState).Start();
     }
diff --git a/src/Wolverine/Runtime/Handlers/HandlerChain.cs b/src/Wolverine/Runtime/Handlers/HandlerChain.cs
index d08e5321b..5937bf052 100644
--- a/src/Wolverine/Runtime/Handlers/HandlerChain.cs
+++ b/src/Wolverine/Runtime/Handlers/HandlerChain.cs
@@ -65,6 +65,18 @@ private HandlerChain(MethodCall call, HandlerGraph parent) : this(call.Method.Me
         Handlers.Add(call);
     }
 
+    internal HandlerChain(MethodCall call, HandlerGraph parent, Endpoint[] endpoints) : this(call, parent)
+    {
+        foreach (var endpoint in endpoints)
+        {
+            RegisterEndpoint(endpoint);
+        }
+        
+        TypeName = call.HandlerType.ToSuffixedTypeName(HandlerSuffix).Replace("[]", "Array");
+
+        Description = $"Message Handler for {MessageType.FullNameInCode()} using {call}";
+    }
+
     public HandlerChain(WolverineOptions options, IGrouping<Type, HandlerCall> grouping, HandlerGraph parent) : this(grouping.Key, parent)
     {
         Handlers.AddRange(grouping);
@@ -111,11 +123,7 @@ private void tryAssignStickyEndpoints(HandlerCall handlerCall, WolverineOptions
                 stub.Subscriptions.Add(Subscription.ForType(MessageType));
             }
             
-            var chain = new HandlerChain(handlerCall, options.HandlerGraph);
-            foreach (var endpoint in endpoints)
-            {
-                chain.RegisterEndpoint(endpoint);
-            }
+            var chain = new HandlerChain(handlerCall, options.HandlerGraph, endpoints);
 
             Handlers.Remove(handlerCall);
             
@@ -199,7 +207,7 @@ public LogLevel ExecutionLogLevel
     /// <summary>
     ///     Wolverine's string identification for this message type
     /// </summary>
-    public string TypeName { get; }
+    public string TypeName { get; private set; }
 
     internal MessageHandler? Handler { get; private set; }
 
diff --git a/src/Wolverine/Runtime/Handlers/HandlerGraph.GeneratesCode.cs b/src/Wolverine/Runtime/Handlers/HandlerGraph.GeneratesCode.cs
index 7a0135c07..fbc0b6dc9 100644
--- a/src/Wolverine/Runtime/Handlers/HandlerGraph.GeneratesCode.cs
+++ b/src/Wolverine/Runtime/Handlers/HandlerGraph.GeneratesCode.cs
@@ -10,6 +10,25 @@ public partial class HandlerGraph
 
     IReadOnlyList<ICodeFile> ICodeFileCollection.BuildFiles()
     {
-        return Chains.ToList();
+
+        return explodeAllFiles().ToList();
+    }
+
+    private IEnumerable<ICodeFile> explodeAllFiles()
+    {
+        foreach (var chain in Chains)
+        {
+            if (chain.Handlers.Any())
+            {
+                yield return chain;
+            }
+            else
+            {
+                foreach (var handlerChain in chain.ByEndpoint)
+                {
+                    yield return handlerChain;
+                }
+            }
+        }
     }
 }
\ No newline at end of file