Add entity source generation support for Durable Functions

src/Generators/DurableTaskSourceGenerator.cs

@@ -270,6 +270,15 @@ public static class GeneratedDurableTaskExtensions
                 }
             }
 
+            foreach (DurableTaskTypeInfo entity in entities)
+            {
+                if (isDurableFunctions)
+                {
+                    // Generate the function definition required to trigger entities in Azure Functions
+                    AddEntityFunctionDeclaration(sourceBuilder, entity);
+                }
+            }
+
             // Activity function triggers are supported for code-gen (but not orchestration triggers)
             IEnumerable<DurableFunction> activityTriggers = allFunctions.Where(
                 df => df.Kind == DurableFunctionKind.Activity);
@@ -368,6 +377,17 @@ static void AddActivityFunctionDeclaration(StringBuilder sourceBuilder, DurableT
         }}");
         }
 
+        static void AddEntityFunctionDeclaration(StringBuilder sourceBuilder, DurableTaskTypeInfo entity)
+        {
+            // Generate the entity trigger function that dispatches to the entity implementation.
+            sourceBuilder.AppendLine($@"
+        [Function(nameof({entity.TaskName}))]
+        public static Task {entity.TaskName}([EntityTrigger] TaskEntityDispatcher dispatcher)
+        {{
+            return dispatcher.DispatchAsync<{entity.TypeName}>();
+        }}");
+        }
+
         /// <summary>
         /// Adds a custom ITaskActivityContext implementation used by code generated from <see cref="AddActivityFunctionDeclaration"/>.
         /// </summary>

test/Generators.Tests/AzureFunctionsTests.cs

@@ -312,4 +312,227 @@ await TestHelpers.RunTestAsync<DurableTaskSourceGenerator>(
             expectedOutput,
             isDurableFunctions: true);
     }
+
+    /// <summary>
+    /// Verifies that using the class-based syntax for authoring entities generates
+    /// <see cref="EntityTriggerAttribute"/> function triggers for Azure Functions.
+    /// </summary>
+    /// <param name="stateType">The entity state type.</param>
+    [Theory]
+    [InlineData("int")]
+    [InlineData("string")]
+    public async Task Entities_ClassBasedSyntax(string stateType)
+    {
+        string code = $@"
+#nullable enable
+using System.Threading.Tasks;
+using Microsoft.DurableTask;
+using Microsoft.DurableTask.Entities;
+
+namespace MyNS
+{{
+    [DurableTask(nameof(MyEntity))]
+    public class MyEntity : TaskEntity<{stateType}>
+    {{
+        public {stateType} Get() => this.State;
+    }}
+}}";
+
+        string expectedOutput = TestHelpers.WrapAndFormat(
+            GeneratedClassName,
+            methodList: @"
+[Function(nameof(MyEntity))]
+public static Task MyEntity([EntityTrigger] TaskEntityDispatcher dispatcher)
+{
+    return dispatcher.DispatchAsync<MyNS.MyEntity>();
+}",
+            isDurableFunctions: true);
+
+        await TestHelpers.RunTestAsync<DurableTaskSourceGenerator>(
+            GeneratedFileName,
+            code,
+            expectedOutput,
+            isDurableFunctions: true);
+    }
+
+    /// <summary>
+    /// Verifies that using the class-based syntax for authoring entities with inheritance generates
+    /// <see cref="EntityTriggerAttribute"/> function triggers for Azure Functions.
+    /// </summary>
+    /// <param name="stateType">The entity state type.</param>
+    [Theory]
+    [InlineData("int")]
+    [InlineData("string")]
+    public async Task Entities_ClassBasedSyntax_Inheritance(string stateType)
+    {
+        string code = $@"
+#nullable enable
+using System.Threading.Tasks;
+using Microsoft.DurableTask;
+using Microsoft.DurableTask.Entities;
+
+namespace MyNS
+{{
+    [DurableTask]
+    public class MyEntity : MyEntityBase
+    {{
+        public override {stateType} Get() => this.State;
+    }}
+
+    public abstract class MyEntityBase : TaskEntity<{stateType}>
+    {{
+        public abstract {stateType} Get();
+    }}
+}}";
+
+        string expectedOutput = TestHelpers.WrapAndFormat(
+            GeneratedClassName,
+            methodList: @"
+[Function(nameof(MyEntity))]
+public static Task MyEntity([EntityTrigger] TaskEntityDispatcher dispatcher)
+{
+    return dispatcher.DispatchAsync<MyNS.MyEntity>();
+}",
+            isDurableFunctions: true);
+
+        await TestHelpers.RunTestAsync<DurableTaskSourceGenerator>(
+            GeneratedFileName,
+            code,
+            expectedOutput,
+            isDurableFunctions: true);
+    }
+
+    /// <summary>
+    /// Verifies that using the class-based syntax for authoring entities with custom state types generates
+    /// <see cref="EntityTriggerAttribute"/> function triggers for Azure Functions.
+    /// </summary>
+    [Fact]
+    public async Task Entities_ClassBasedSyntax_CustomStateType()
+    {
+        string code = @"
+#nullable enable
+using System.Threading.Tasks;
+using Microsoft.DurableTask;
+using Microsoft.DurableTask.Entities;
+
+namespace MyNS
+{
+    public class MyState
+    {
+        public int Value { get; set; }
+    }
+
+    [DurableTask(nameof(MyEntity))]
+    public class MyEntity : TaskEntity<MyState>
+    {
+        public MyState Get() => this.State;
+    }
+}";
+
+        string expectedOutput = TestHelpers.WrapAndFormat(
+            GeneratedClassName,
+            methodList: @"
+[Function(nameof(MyEntity))]
+public static Task MyEntity([EntityTrigger] TaskEntityDispatcher dispatcher)
+{
+    return dispatcher.DispatchAsync<MyNS.MyEntity>();
+}",
+            isDurableFunctions: true);
+
+        await TestHelpers.RunTestAsync<DurableTaskSourceGenerator>(
+            GeneratedFileName,
+            code,
+            expectedOutput,
+            isDurableFunctions: true);
+    }
+
+    /// <summary>
+    /// Verifies that using the class-based syntax for authoring a mix of orchestrators, activities,
+    /// and entities generates the appropriate function triggers for Azure Functions.
+    /// </summary>
+    [Fact]
+    public async Task Mixed_OrchestratorActivityEntity_ClassBasedSyntax()
+    {
+        string code = @"
+#nullable enable
+using System;
+using System.Threading.Tasks;
+using Microsoft.DurableTask;
+using Microsoft.DurableTask.Entities;
+
+namespace MyNS
+{
+    [DurableTask(nameof(MyOrchestrator))]
+    public class MyOrchestrator : TaskOrchestrator<int, string>
+    {
+        public override Task<string> RunAsync(TaskOrchestrationContext ctx, int input) => Task.FromResult(string.Empty);
+    }
+
+    [DurableTask(nameof(MyActivity))]
+    public class MyActivity : TaskActivity<int, string>
+    {
+        public override Task<string> RunAsync(TaskActivityContext context, int input) => Task.FromResult(string.Empty);
+    }
+
+    [DurableTask(nameof(MyEntity))]
+    public class MyEntity : TaskEntity<int>
+    {
+        public int Get() => this.State;
+    }
+}";
+
+        string expectedOutput = TestHelpers.WrapAndFormat(
+            GeneratedClassName,
+            methodList: $@"
+static readonly ITaskOrchestrator singletonMyOrchestrator = new MyNS.MyOrchestrator();
+
+[Function(nameof(MyOrchestrator))]
+public static Task<string> MyOrchestrator([OrchestrationTrigger] TaskOrchestrationContext context)
+{{
+    return singletonMyOrchestrator.RunAsync(context, context.GetInput<int>())
+        .ContinueWith(t => (string)(t.Result ?? default(string)!), TaskContinuationOptions.ExecuteSynchronously);
+}}
+
+/// <inheritdoc cref=""IOrchestrationSubmitter.ScheduleNewOrchestrationInstanceAsync""/>
+public static Task<string> ScheduleNewMyOrchestratorInstanceAsync(
+    this IOrchestrationSubmitter client, int input, StartOrchestrationOptions? options = null)
+{{
+    return client.ScheduleNewOrchestrationInstanceAsync(""MyOrchestrator"", input, options);
+}}
+
+/// <inheritdoc cref=""TaskOrchestrationContext.CallSubOrchestratorAsync(TaskName, object?, TaskOptions?)""/>
+public static Task<string> CallMyOrchestratorAsync(
+    this TaskOrchestrationContext context, int input, TaskOptions? options = null)
+{{
+    return context.CallSubOrchestratorAsync<string>(""MyOrchestrator"", input, options);
+}}
+
+public static Task<string> CallMyActivityAsync(this TaskOrchestrationContext ctx, int input, TaskOptions? options = null)
+{{
+    return ctx.CallActivityAsync<string>(""MyActivity"", input, options);
+}}
+
+[Function(nameof(MyActivity))]
+public static async Task<string> MyActivity([ActivityTrigger] int input, string instanceId, FunctionContext executionContext)
+{{
+    ITaskActivity activity = ActivatorUtilities.GetServiceOrCreateInstance<MyNS.MyActivity>(executionContext.InstanceServices);
+    TaskActivityContext context = new GeneratedActivityContext(""MyActivity"", instanceId);
+    object? result = await activity.RunAsync(context, input);
+    return (string)result!;
+}}
+
+[Function(nameof(MyEntity))]
+public static Task MyEntity([EntityTrigger] TaskEntityDispatcher dispatcher)
+{{
+    return dispatcher.DispatchAsync<MyNS.MyEntity>();
+}}
+{TestHelpers.DeIndent(DurableTaskSourceGenerator.GetGeneratedActivityContextCode(), spacesToRemove: 8)}",
+            isDurableFunctions: true);
+
+        await TestHelpers.RunTestAsync<DurableTaskSourceGenerator>(
+            GeneratedFileName,
+            code,
+            expectedOutput,
+            isDurableFunctions: true);
+    }
 }