Add visitor for MultiPartFormData WriteTo method (Azure#52933)

* Add visitor for MultiPartFormData WriteTo method

* Include TaskExtensions

* fix

* revert

* regen

Author: JoshLove-msft

eng/packages/http-client-csharp/generator/Azure.Generator/src/AzureClientGenerator.cs

@@ -64,5 +64,6 @@ protected override void Configure()
         AddVisitor(new MatchConditionsHeadersVisitor());
         AddVisitor(new RequestClientIdHeaderVisitor());
         AddVisitor(new SystemTextJsonConverterVisitor());
+        AddVisitor(new MultiPartFormDataVisitor());
     }
 }

eng/packages/http-client-csharp/generator/Azure.Generator/src/Primitives/NewAzureProjectScaffolding.cs

@@ -173,6 +173,12 @@ protected override IReadOnlyList<CSharpProjectCompileInclude> BuildCompileInclud
                 }
             }
 
+            // Add TaskExtensions if there are multipart form data operations and it hasn't already been added for LRO
+            if (!hasLongRunningOperation && AzureClientGenerator.Instance.InputLibrary.HasMultipartFormDataOperation)
+            {
+                compileIncludes.Add(new CSharpProjectCompileInclude(GetCompileInclude("TaskExtensions.cs"), SharedSourceLinkBase));
+            }
+
             return compileIncludes;
         }
     }

eng/packages/http-client-csharp/generator/Azure.Generator/src/Visitors/MultiPartFormDataVisitor.cs

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Azure.Core.Pipeline;
+using Microsoft.TypeSpec.Generator.ClientModel;
+using Microsoft.TypeSpec.Generator.ClientModel.Providers;
+using Microsoft.TypeSpec.Generator.Expressions;
+using Microsoft.TypeSpec.Generator.Providers;
+using Microsoft.TypeSpec.Generator.Statements;
+
+namespace Azure.Generator.Visitors
+{
+    internal class MultiPartFormDataVisitor : ScmLibraryVisitor
+    {
+        protected override MethodBodyStatement VisitExpressionStatement(ExpressionStatement statement, MethodProvider method)
+        {
+            if (method is { EnclosingType: MultiPartFormDataBinaryContentDefinition, Signature.Name: "WriteTo" })
+            {
+                // Replaces GetAwaiter().GetResult() with EnsureCompleted() for the sync over async path.
+                if (statement.Expression is InvokeMethodExpression invokeMethod)
+                {
+                    if (invokeMethod.InstanceReference is InvokeMethodExpression innerInvoke)
+                    {
+                        invokeMethod.Update(
+                            instanceReference: innerInvoke.InstanceReference,
+                            methodName: nameof(TaskExtensions.EnsureCompleted),
+                            extensionType: typeof(TaskExtensions));
+                    }
+                }
+            }
+
+            return statement;
+        }
+    }
+}

eng/packages/http-client-csharp/generator/Azure.Generator/test/Visitors/MultiPartFormDataVisitorTests.cs

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System.Linq;
+using Azure.Generator.Tests.Common;
+using Azure.Generator.Tests.TestHelpers;
+using Azure.Generator.Visitors;
+using Microsoft.TypeSpec.Generator;
+using Microsoft.TypeSpec.Generator.ClientModel.Providers;
+using NUnit.Framework;
+
+namespace Azure.Generator.Tests.Visitors
+{
+    public class MultiPartFormDataVisitorTests
+    {
+        [Test]
+        public void UpdatesWriteToMethod()
+        {
+            var operation = InputFactory.Operation("SomeOperation", requestMediaTypes: ["multipart/form-data"]);
+            var serviceMethod = InputFactory.BasicServiceMethod("SomeOperation", operation);
+            var client = InputFactory.Client("client", methods: [serviceMethod]);
+            var plugin = MockHelpers.LoadMockGenerator(clients: () => [client]);
+
+            var visitor = new TestMultiPartFormDataVisitor();
+            visitor.InvokeVisitLibrary(plugin.Object.OutputLibrary);
+
+            var multiPartFormData = plugin.Object.OutputLibrary.TypeProviders.OfType<MultiPartFormDataBinaryContentDefinition>().SingleOrDefault();
+
+            Assert.IsNotNull(multiPartFormData);
+            var writeToMethod = multiPartFormData!.Methods.SingleOrDefault(m => m.Signature.Name == "WriteTo");
+            Assert.IsNotNull(writeToMethod);
+            Assert.AreEqual(Helpers.GetExpectedFromFile(), writeToMethod!.BodyStatements!.ToDisplayString());
+        }
+
+        private class TestMultiPartFormDataVisitor : MultiPartFormDataVisitor
+        {
+            public void InvokeVisitLibrary(OutputLibrary library)
+            {
+                base.VisitLibrary(library);
+            }
+        }
+    }
+}

eng/packages/http-client-csharp/generator/Azure.Generator/test/Visitors/TestData/MultiPartFormDataVisitorTests/UpdatesWriteToMethod.cs

@@ -0,0 +1,5 @@
+#if NET6_0_OR_GREATER
+_multipartContent.CopyTo(stream, default, cancellationToken);
+#else
+_multipartContent.CopyToAsync(stream).EnsureCompleted();
+#endif

eng/packages/http-client-csharp/generator/TestProjects/Local/Basic-TypeSpec/Basic-TypeSpec.tsp

@@ -450,4 +450,10 @@ op ConditionalRequest(@header("If-Match") ifMatch?: string, @header("If-None-Mat
 @doc("A sample operation with conditional requests")
 @convenientAPI(true)
 @get
-op ConditionalRequestDate(@header("If-Match") ifMatch?: string, @header("If-None-Match") ifNoneMatch?: string, @header("If-Modified-Since") ifModifiedSince?: string): void;
+op ConditionalRequestDate(@header("If-Match") ifMatch?: string, @header("If-None-Match") ifNoneMatch?: string, @header("If-Modified-Since") ifModifiedSince?: string): void;
+
+@route("multipart")
+@doc("A sample operation with multipart request")
+@convenientAPI(true)
+@post
+op MultipartRequest(@header("Content-Type") contentType: "multipart/form-data", @multipartBody body: {name: HttpPart<string>}): void;

eng/packages/http-client-csharp/generator/TestProjects/Local/Basic-TypeSpec/src/BasicTypeSpec.csproj

@@ -18,6 +18,7 @@
     <Compile Include="$(AzureCoreSharedSources)DiagnosticScope.cs" LinkBase="Shared/Core" />
     <Compile Include="$(AzureCoreSharedSources)HttpMessageSanitizer.cs" LinkBase="Shared/Core" />
     <Compile Include="$(AzureCoreSharedSources)TrimmingAttribute.cs" LinkBase="Shared/Core" />
+    <Compile Include="$(AzureCoreSharedSources)TaskExtensions.cs" LinkBase="Shared/Core" />
   </ItemGroup>
 
   <ItemGroup>