Analyzer fix for handlers declared on different partial classes (#6372)

* Analyzer fix for handlers declared on different partial classes

* Update src/NServiceBus.Core.Analyzer.Tests.Common/Sagas/SagaAnalyzerTests.cs

Author: DavidBoike

src/NServiceBus.Core.Analyzer.Tests.Common/Sagas/SagaAnalyzerTests.cs

@@ -1028,6 +1028,115 @@ public Task Handle(StartSaga message, IMessageHandlerContext context)
                 return Assert(source);
             }
         }
+
+        // https://github.com/Particular/NServiceBus/issues/6370
+        [Test]
+        public Task SagaHandlersInPartialClasses()
+        {
+            var source =
+@"
+using System.Threading.Tasks;
+using NServiceBus;
+
+public partial class SagaImplementation : Saga<SagaData>, IAmStartedByMessages<SagaStartMessage>
+{
+    protected override void ConfigureHowToFindSaga(SagaPropertyMapper<SagaData> mapper)
+    {
+        mapper.MapSaga(s => s.OrderId)
+            .ToMessage<SagaStartMessage>(b => b.OrderId)
+            .ToMessage<SagaStep1>(b => b.OrderId)
+            .ToMessage<SagaStep2>(s => s.OrderId)
+            .ToMessage<SagaStep3>(b => b.OrderId);
+    }
+
+    public async Task Handle(SagaStartMessage message, IMessageHandlerContext context)
+    {
+        var options = new SendOptions();
+        options.RouteToThisEndpoint();
+        await context.Send(new SagaStep1 {OrderId = Data.OrderId}, options);
+        await context.Send(new SagaStep2 {OrderId = Data.OrderId}, options);
+        await context.Send(new SagaStep3 {OrderId = Data.OrderId}, options);
+    }
+
+    private void CompleteSaga()
+    {
+        if(Data.Step1Complete && Data.Step2Complete && Data.Step3Complete)
+        {
+            MarkAsComplete();
+        }
+    }
+}
+-----
+using System.Threading.Tasks;
+using NServiceBus;
+
+public partial class SagaImplementation : IHandleMessages<SagaStep1>
+{
+    public Task Handle(SagaStep1 message, IMessageHandlerContext context)
+    {
+        Data.Step1Complete = true;
+        CompleteSaga();
+        return Task.CompletedTask;
+    }
+}
+-----
+using System.Threading.Tasks;
+using NServiceBus;
+
+public partial class SagaImplementation : IHandleMessages<SagaStep2>
+{
+    public  Task Handle(SagaStep2 message, IMessageHandlerContext context)
+    {
+        Data.Step2Complete = true;
+        CompleteSaga();
+        return Task.CompletedTask;
+    }
+}
+-----
+using System.Threading.Tasks;
+using NServiceBus;
+
+public partial class SagaImplementation : IHandleMessages<SagaStep3>
+{
+    public Task Handle(SagaStep3 message, IMessageHandlerContext context)
+    {
+        Data.Step3Complete = true;
+        CompleteSaga();
+        return Task.CompletedTask;
+    }
+}
+-----
+using NServiceBus;
+
+public class SagaData : ContainSagaData
+{
+    public string OrderId { get; set; } = null!;
+
+    public bool Step1Complete { get; set; }
+    public bool Step2Complete { get; set; }
+    public bool Step3Complete { get; set; }
+}
+
+public class SagaStartMessage : ICommand
+{
+    public string OrderId { get; set; } = null!;
+}
+
+public class SagaStep1 : SagaStartMessage
+{
+}
+
+public class SagaStep2 : SagaStartMessage
+{
+}
+
+public class SagaStep3 : SagaStartMessage
+{
+}
+";
+
+            return Assert(source);
+        }
     }
 
     public class SagaAnalyzerTestsCSharp9 : SagaAnalyzerTestsCSharp8

src/NServiceBus.Core.Analyzer/Sagas/SagaAnalyzer.cs

@@ -102,7 +102,10 @@ static void AnalyzeSagaClass(SyntaxNodeAnalysisContext context, ClassDeclaration
                 }
             }
 
-            if (!TryGetSagaDetails(context, knownTypes, sagaType, out var saga))
+            // Manages access to semantic models for other files, since those are expensive to create
+            var semanticModels = new SemanticModelCache(context.Compilation, context.Node.SyntaxTree, context.SemanticModel);
+
+            if (!TryGetSagaDetails(context, semanticModels, knownTypes, sagaType, out var saga))
             {
                 return;
             }
@@ -235,7 +238,11 @@ static void AnalyzeSagaClass(SyntaxNodeAnalysisContext context, ClassDeclaration
 
             // Figure out which message types have a ConfigureHowToFindSaga mapping...
             var mappedMessageTypes = saga.MessageMappings
-                .Select(m => context.SemanticModel.GetTypeInfo(m.MessageTypeSyntax).Type)
+                .Select(m =>
+                {
+                    var semanticModel = semanticModels.GetFor(m.MessageTypeSyntax);
+                    return semanticModel.GetTypeInfo(m.MessageTypeSyntax).Type;
+                })
                 .ToImmutableHashSet(SymbolEqualityComparer.Default); // Message types shouldn't need nullability annotations
 
             // ...then find the IAmStartedBy message types that don't already have a mapping defined
@@ -325,7 +332,10 @@ static void AnalyzeSagaDataClass(SyntaxNodeAnalysisContext context, ClassDeclara
                 return;
             }
 
-            if (!TryGetSagaDetails(context, knownTypes, sagaType, out var saga))
+            // Manages access to semantic models for other files, since those are expensive to create
+            var semanticModels = new SemanticModelCache(context.Compilation, context.Node.SyntaxTree, context.SemanticModel);
+
+            if (!TryGetSagaDetails(context, semanticModels, knownTypes, sagaType, out var saga))
             {
                 return;
             }
@@ -487,7 +497,7 @@ static Diagnostic CreateMappingRewritingDiagnostic(
             return diagnostic;
         }
 
-        static bool TryGetSagaDetails(SyntaxNodeAnalysisContext context, KnownTypes knownTypes, INamedTypeSymbol sagaType, out SagaDetails saga)
+        static bool TryGetSagaDetails(SyntaxNodeAnalysisContext context, SemanticModelCache semanticModels, KnownTypes knownTypes, INamedTypeSymbol sagaType, out SagaDetails saga)
         {
             saga = null;
 
@@ -515,9 +525,6 @@ static bool TryGetSagaDetails(SyntaxNodeAnalysisContext context, KnownTypes know
                 return false;
             }
 
-            // Manages access to semantic models for other files, since those are expensive to create
-            var semanticModels = new SemanticModelCache(context.Compilation, context.Node.SyntaxTree, context.SemanticModel);
-
             // From all the base lists on all partials, find the methods that are one of our IAmStarted/IHandle methods, then get
             // the TypeSymbol so we have both the syntax and type available
             var handlerDeclarations = classDeclarations