Fix CancellationTokenAnalyzer fluent style API bug

src/NServiceBus.Core.Analyzer.Tests.Common/ForwardCancellationToken/ForwardCancellationTokenFixerTests.cs

@@ -42,6 +42,96 @@ public Task Bar(IMessageHandlerContext context)
             return Assert(original, expected);
         }
 
+        [Test]
+        public Task Fluent()
+        {
+            var original =
+                @"using NServiceBus;
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+public class Foo
+{
+    public Task Bar(IMessageHandlerContext context)
+    {
+        var someReturnValue = ""someValue"";
+        return 
+             Include(() => someReturnValue)
+            .Include(() => someReturnValue)
+            .FindSingle();
+    }
+
+    Foo Include(Func<string> action) => this;
+
+    Task FindSingle(CancellationToken token = default(CancellationToken)) { return Task.CompletedTask; }
+}";
+
+            var expected =
+                @"using NServiceBus;
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+public class Foo
+{
+    public Task Bar(IMessageHandlerContext context)
+    {
+        var someReturnValue = ""someValue"";
+        return 
+             Include(() => someReturnValue)
+            .Include(() => someReturnValue)
+            .FindSingle(context.CancellationToken);
+    }
+
+    Foo Include(Func<string> action) => this;
+
+    Task FindSingle(CancellationToken token = default(CancellationToken)) { return Task.CompletedTask; }
+}";
+
+            return Assert(original, expected);
+        }
+
+        [Test]
+        public Task FluentAsync()
+        {
+            var original =
+                @"using NServiceBus;
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+public class Foo
+{
+    public async Task Bar(IMessageHandlerContext context)
+    {
+        var someReturnValue = ""someValue"";
+        await (await (await Include(() => someReturnValue)).Include(() => someReturnValue)).FindSingle();
+    }
+
+    Task<Foo> Include(Func<string> action, CancellationToken token = default(CancellationToken)) => Task.FromResult(this);
+
+    Task<Foo> FindSingle(CancellationToken token = default(CancellationToken)) => Task.FromResult(this);
+}";
+
+            var expected =
+                @"using NServiceBus;
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+public class Foo
+{
+    public async Task Bar(IMessageHandlerContext context)
+    {
+        var someReturnValue = ""someValue"";
+        await (await (await Include(() => someReturnValue, context.CancellationToken)).Include(() => someReturnValue, context.CancellationToken)).FindSingle(context.CancellationToken);
+    }
+
+    Task<Foo> Include(Func<string> action, CancellationToken token = default(CancellationToken)) => Task.FromResult(this);
+
+    Task<Foo> FindSingle(CancellationToken token = default(CancellationToken)) => Task.FromResult(this);
+}";
+
+            return Assert(original, expected);
+        }
+
         [Test]
         public Task NonStandardContextVariableName()
         {
@@ -245,4 +335,4 @@ public class ForwardCancellationTokenFixerTestsCSharp10 : ForwardCancellationTok
         protected override LanguageVersion AnalyzerLanguageVersion => LanguageVersion.CSharp10;
     }
 #endif
-}
+}

src/NServiceBus.Core.Analyzer/ForwardCancellationTokenFixer.cs

@@ -34,8 +34,8 @@ public sealed override async Task RegisterCodeFixesAsync(CodeFixContext context)
             var requiredArgName = diagnostic.Properties["RequiredArgName"];
 
             var root = await context.Document.GetSyntaxRootAsync(context.CancellationToken).ConfigureAwait(false);
-            var startToken = root.FindToken(diagnostic.Location.SourceSpan.Start);
-            var invocationSyntax = startToken.Parent.AncestorsAndSelf().OfType<InvocationExpressionSyntax>().First();
+            var startToken = root.FindNode(diagnostic.Location.SourceSpan);
+            var invocationSyntax = (InvocationExpressionSyntax)startToken;
 
             // TODO: consider whether analyzer should pass in name of property, i.e. "CancellationToken",
             // so that that information doesn't have to spread across both the analyzer and the fixer
@@ -89,4 +89,4 @@ static async Task<Document> ForwardCancellationToken(
         // the title is dynamic based on the context name and target method name.
         static readonly string EquivalenceKey = typeof(ForwardCancellationTokenFixer).FullName;
     }
-}
+}