ExceptionHandlingIntegrationTests.cs

//  ----------------------------------------------------------------------------------
//  Copyright Microsoft Corporation
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//  http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//  ----------------------------------------------------------------------------------
#nullable enable
namespace DurableTask.Core.Tests
{
    using System;
    using System.Collections.Generic;
    using System.Diagnostics;
    using System.Reflection;
    using System.Reflection.Emit;
    using System.Runtime.Serialization;
    using System.Threading.Tasks;
    using DurableTask.Core.Exceptions;
    using DurableTask.Emulator;
    using Microsoft.Extensions.Logging;
    using Microsoft.VisualStudio.TestTools.UnitTesting;
    using Newtonsoft.Json;

    [TestClass]
    public class ExceptionHandlingIntegrationTests
    {
        static readonly TimeSpan DefaultTimeout = TimeSpan.FromSeconds(Debugger.IsAttached ? 300 : 10);

        readonly TaskHubWorker worker;
        readonly TaskHubClient client;

        public ExceptionHandlingIntegrationTests()
        {
            // configure logging so traces are emitted during tests.
            // This facilitates debugging when tests fail.

            var loggerFactory = LoggerFactory.Create(builder =>
            {
                builder.AddConsole().SetMinimumLevel(LogLevel.Trace);
            });

            var service = new LocalOrchestrationService();
            this.worker = new TaskHubWorker(service, loggerFactory);
            this.client = new TaskHubClient(service, loggerFactory: loggerFactory);
        }

        [DataTestMethod]
        [DataRow(ErrorPropagationMode.SerializeExceptions)]
        [DataRow(ErrorPropagationMode.UseFailureDetails)]
        public async Task CatchInvalidOperationException(ErrorPropagationMode mode)
        {
            // The error propagation mode must be set before the worker is started
            this.worker.ErrorPropagationMode = mode;

            await this.worker
                .AddTaskOrchestrations(typeof(ExceptionHandlingOrchestration))
                .AddTaskActivities(typeof(ThrowInvalidOperationException))
                .StartAsync();

            // This is required for exceptions to be serialized
            this.worker.TaskActivityDispatcher.IncludeDetails = true;

            OrchestrationInstance instance = await this.client.CreateOrchestrationInstanceAsync(typeof(ExceptionHandlingOrchestration), null);
            OrchestrationState state = await this.client.WaitForOrchestrationAsync(instance, DefaultTimeout);
            Assert.IsNotNull(state);
            Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus);
            Assert.IsNotNull(state.Output, "The expected error information wasn't found!");

            if (mode == ErrorPropagationMode.SerializeExceptions)
            {
                // The exception should be deserializable
                InvalidOperationException? e = JsonConvert.DeserializeObject<InvalidOperationException>(state.Output);
                Assert.IsNotNull(e);
                Assert.AreEqual("This is a test exception", e!.Message);
            }
            else if (mode == ErrorPropagationMode.UseFailureDetails)
            {
                // The failure details should contain the relevant exception metadata
                FailureDetails? details = JsonConvert.DeserializeObject<FailureDetails>(state.Output);
                Assert.IsNotNull(details);
                Assert.AreEqual(typeof(InvalidOperationException).FullName, details!.ErrorType);
                Assert.IsTrue(details.IsCausedBy<InvalidOperationException>());
                Assert.IsTrue(details.IsCausedBy<Exception>()); // check that base types work too
                Assert.AreEqual("This is a test exception", details.ErrorMessage);
                Assert.IsNotNull(details.StackTrace);

                // The callstack should be in the error details
                string expectedCallstackSubstring = typeof(ThrowInvalidOperationException).FullName!.Replace('+', '.');
                Assert.IsTrue(
                    details.StackTrace!.IndexOf(expectedCallstackSubstring) > 0,
                    $"Expected to find {expectedCallstackSubstring} in the exception details. Actual: {details.StackTrace}");
            }
            else
            {
                Assert.Fail($"Unexpected {nameof(ErrorPropagationMode)} value: {mode}");
            }
        }

        [TestMethod]
        public async Task FailureDetailsOnHandled()
        {
            // The error propagation mode must be set before the worker is started
            this.worker.ErrorPropagationMode = ErrorPropagationMode.UseFailureDetails;

            await this.worker
                .AddTaskOrchestrations(typeof(ExceptionHandlingWithRetryOrchestration))
                .AddTaskActivities(typeof(ThrowInvalidOperationException))
                .StartAsync();

            OrchestrationInstance instance = await this.client.CreateOrchestrationInstanceAsync(typeof(ExceptionHandlingWithRetryOrchestration), null);
            OrchestrationState state = await this.client.WaitForOrchestrationAsync(instance, DefaultTimeout);
            Assert.IsNotNull(state);
            Assert.AreEqual(OrchestrationStatus.Completed, state.OrchestrationStatus);
            Assert.IsNotNull(state.Output, "No output was returned!");

            // The orchestration is written in such a way that there should be only one call into the retry policy
            int retryPolicyInvokedCount = JsonConvert.DeserializeObject<int>(state.Output);
            Assert.AreEqual(1, retryPolicyInvokedCount);
        }

        [DataTestMethod]
        [DataRow(ErrorPropagationMode.SerializeExceptions)]
        [DataRow(ErrorPropagationMode.UseFailureDetails)]
        public async Task FailureDetailsOnUnhandled(ErrorPropagationMode mode)
        {
            // The error propagation mode must be set before the worker is started
            this.worker.ErrorPropagationMode = mode;

            await this.worker
                .AddTaskOrchestrations(typeof(NoExceptionHandlingOrchestration))
                .AddTaskActivities(typeof(ThrowInvalidOperationException))
                .StartAsync();

            OrchestrationInstance instance = await this.client.CreateOrchestrationInstanceAsync(
                typeof(NoExceptionHandlingOrchestration),
                input: null);
            OrchestrationState state = await this.client.WaitForOrchestrationAsync(instance, DefaultTimeout);
            Assert.IsNotNull(state);
            Assert.AreEqual(OrchestrationStatus.Failed, state.OrchestrationStatus);

            string expectedErrorMessage = "This is a test exception";
            if (mode == ErrorPropagationMode.SerializeExceptions)
            {
                // Legacy behavior is to set the output of the orchestration to be the exception message
                Assert.AreEqual(expectedErrorMessage, state.Output);
            }
            else if (mode == ErrorPropagationMode.UseFailureDetails)
            {
                string activityName = typeof(ThrowInvalidOperationException).FullName!;
                string expectedOutput = $"{typeof(TaskFailedException).FullName}: Task '{activityName}' (#0) failed with an unhandled exception: {expectedErrorMessage}";
                Assert.AreEqual(expectedOutput, state.Output);
           }
            else
            {
                Assert.Fail($"Unexpected {nameof(ErrorPropagationMode)} value: {mode}");
            }
        }

        [TestMethod]
        public void TaskFailureOnNullContextTaskActivity()
        {
            TaskActivity activity = new ThrowInvalidOperationExceptionAsync();
            string input = JsonConvert.SerializeObject(new string[] { "test" });

            // Pass a null context to check that it doesn't affect error handling.
            Task<string> task = activity.RunAsync(null, input);

            Assert.IsTrue(task.IsFaulted);
            Assert.IsNotNull(task.Exception);
            Assert.IsNotNull(task.Exception?.InnerException);
            Assert.IsInstanceOfType(task.Exception?.InnerException, typeof(TaskFailureException));
            Assert.AreEqual("This is a test exception", task.Exception?.InnerException?.Message);
        }

        class ExceptionHandlingOrchestration : TaskOrchestration<object, string>
        {
            public override async Task<object> RunTask(OrchestrationContext context, string input)
            {
                try
                {
                    return await context.ScheduleTask<object>(typeof(ThrowInvalidOperationException));
                }
                catch (TaskFailedException e)
                {
                    // Exactly one of these properties should be null
                    return (object)e.FailureDetails! ?? e.InnerException!;
                }
            }
        }

        class ExceptionHandlingWithRetryOrchestration : TaskOrchestration<int, string>
        {
            public override async Task<int> RunTask(OrchestrationContext context, string input)
            {
                int handleCount = 0;
                try
                {
                    await context.ScheduleWithRetry<object>(
                        typeof(ThrowInvalidOperationException),
                        new RetryOptions(TimeSpan.FromMilliseconds(1), maxNumberOfAttempts: 3)
                        {
                            Handle = e =>
                            {
                                handleCount++;

                                // Users should be able to examine the structured exception details when
                                // ErrorPropagationMode is set to UseFailureDetails
                                if (e is TaskFailedException tfe &&
                                    tfe.FailureDetails != null &&
                                    tfe.FailureDetails.ErrorType == typeof(InvalidOperationException).FullName &&
                                    tfe.FailureDetails.ErrorMessage == "This is a test exception" &&
                                    tfe.FailureDetails.StackTrace!.Contains(typeof(ThrowInvalidOperationException).Name) &&
                                    tfe.FailureDetails.IsCausedBy<InvalidOperationException>() &&
                                    tfe.FailureDetails.IsCausedBy<Exception>() &&
                                    tfe.FailureDetails.InnerFailure != null &&
                                    tfe.FailureDetails.InnerFailure.IsCausedBy<CustomException>() &&
                                    tfe.FailureDetails.InnerFailure.ErrorMessage == "And this is its custom inner exception") 
                                {
                                    // Stop retrying
                                    return false;
                                }

                                // Keep retrying
                                return true;
                            }
                        });
                }
                catch (TaskFailedException)
                {
                }

                return handleCount;
            }
        }

        class NoExceptionHandlingOrchestration : TaskOrchestration<object, string>
        {
            public override Task<object> RunTask(OrchestrationContext context, string input)
            {
                // let the exception go unhandled and fail the orchestration
                return context.ScheduleTask<object>(typeof(ThrowInvalidOperationException));
            }
        }

        class ThrowInvalidOperationException : TaskActivity<string, string>
        {
            protected override string Execute(TaskContext context, string input)
            {
                throw new InvalidOperationException("This is a test exception",
                    new CustomException("And this is its custom inner exception"));
            }
        }

        class ThrowInvalidOperationExceptionAsync : AsyncTaskActivity<string, string>
        {
            protected override Task<string> ExecuteAsync(TaskContext context, string input)
            {
                throw new InvalidOperationException("This is a test exception",
                    new CustomException("And this is its custom inner exception"));
            }
        }

        [TestMethod]
        // Test that when a provider is set, properties are extracted and stored in FailureDetails.Properties.
        public async Task ExceptionPropertiesProvider_ExtractsCustomProperties()
        {
            // Set up a provider that extracts custom properties using the new TaskHubWorker property
            this.worker.ExceptionPropertiesProvider = new TestExceptionPropertiesProvider();
            this.worker.ErrorPropagationMode = ErrorPropagationMode.UseFailureDetails;

            try
            {
                await this.worker
                    .AddTaskOrchestrations(typeof(ThrowCustomExceptionOrchestration))
                    .AddTaskActivities(typeof(ThrowCustomBusinessExceptionActivity))
                    .StartAsync();

                var instance = await this.client.CreateOrchestrationInstanceAsync(typeof(ThrowCustomExceptionOrchestration), "test-input");
                var result = await this.client.WaitForOrchestrationAsync(instance, DefaultTimeout);

                // Check that custom properties were extracted
                Assert.AreEqual(OrchestrationStatus.Failed, result.OrchestrationStatus);
                Assert.IsNotNull(result.FailureDetails);
                Assert.IsNotNull(result.FailureDetails.Properties);

                // Check the properties match the exception.
                Assert.AreEqual("CustomBusinessException", result.FailureDetails.Properties["ExceptionTypeName"]);
                Assert.AreEqual("user123", result.FailureDetails.Properties["UserId"]);
                Assert.AreEqual("OrderProcessing", result.FailureDetails.Properties["BusinessContext"]);
                Assert.IsTrue(result.FailureDetails.Properties.ContainsKey("Timestamp"));

                // Check that null values are properly handled
                Assert.IsTrue(result.FailureDetails.Properties.ContainsKey("TestNullObject"), "TestNullObject key should be present");
                Assert.IsNull(result.FailureDetails.Properties["TestNullObject"], "TestNullObject should be null");
                
                Assert.IsTrue(result.FailureDetails.Properties.ContainsKey("DirectNullValue"), "DirectNullValue key should be present");
                Assert.IsNull(result.FailureDetails.Properties["DirectNullValue"], "DirectNullValue should be null");

                // Verify non-null values still work
                Assert.IsTrue(result.FailureDetails.Properties.ContainsKey("EmptyString"), "EmptyString key should be present");
                Assert.AreEqual(string.Empty, result.FailureDetails.Properties["EmptyString"], "EmptyString should be empty string, not null");
            }
            finally
            {
                await this.worker.StopAsync();
            }
        }

        [TestMethod]
        // Test that when no provider is provided by default, property at FailureDetails should be null.
        public async Task ExceptionPropertiesProvider_NullProvider_NoProperties()
        {
            try
            {
                this.worker.ErrorPropagationMode = ErrorPropagationMode.UseFailureDetails;
                await this.worker
                    .AddTaskOrchestrations(typeof(ThrowInvalidOperationExceptionOrchestration))
                    .AddTaskActivities(typeof(ThrowInvalidOperationExceptionActivity))
                    .StartAsync();

                var instance = await this.client.CreateOrchestrationInstanceAsync(typeof(ThrowInvalidOperationExceptionOrchestration), "test-input");
                var result = await this.client.WaitForOrchestrationAsync(instance, DefaultTimeout);

                // Properties should be null when no provider
                Assert.AreEqual(OrchestrationStatus.Failed, result.OrchestrationStatus);
                Assert.IsNotNull(result.FailureDetails);
                Assert.IsNull(result.FailureDetails.Properties);
            }
            finally
            {
                await this.worker.StopAsync();
            }
        }

        [TestMethod]
        // Test that when a provider is set, properties of exception thrown by orchestration directly will be included
        // if excception type is matched.
        public async Task ExceptionPropertiesProvider_SimpleThrowExceptionOrchestration()
        {
            this.worker.ExceptionPropertiesProvider = new TestExceptionPropertiesProvider();
            this.worker.ErrorPropagationMode = ErrorPropagationMode.UseFailureDetails;

            try
            {
                await this.worker
                    .AddTaskOrchestrations(typeof(SimpleThrowExceptionOrchestration))
                    .StartAsync();

                var instance = await this.client.CreateOrchestrationInstanceAsync(typeof(SimpleThrowExceptionOrchestration), "test-input");
                var result = await this.client.WaitForOrchestrationAsync(instance, DefaultTimeout);

                // Check that custom properties were extracted
                Assert.AreEqual(OrchestrationStatus.Failed, result.OrchestrationStatus);
                Assert.IsNotNull(result.FailureDetails);
                Assert.IsNotNull(result.FailureDetails.Properties);

                // Check the properties match the ArgumentOutOfRangeException.
                Assert.AreEqual("count", result.FailureDetails.Properties["Name"]);
                Assert.AreEqual("100", result.FailureDetails.Properties["Value"]);
            }
            finally
            {
                await this.worker.StopAsync();
            }
        }

        [TestMethod]
        // Test that when a provider is set, exception properties are included in failure details with propogation.
        public async Task ExceptionPropertiesProvider_SubOrchestrationThrowExceptionOrchestration()
        {
            this.worker.ExceptionPropertiesProvider = new TestExceptionPropertiesProvider();
            this.worker.ErrorPropagationMode = ErrorPropagationMode.UseFailureDetails;

            try
            {
                await this.worker
                    .AddTaskOrchestrations(typeof(SubOrchestrationThrowExceptionOrchestration))
                    .AddTaskOrchestrations(typeof(ThrowArgumentOutofRangeExceptionASubOrchestration))
                    .AddTaskActivities(typeof(ThrowArgumentOutofRangeExceptionActivity))
                    .StartAsync();

                var instance = await this.client.CreateOrchestrationInstanceAsync(typeof(SubOrchestrationThrowExceptionOrchestration), "test-input");
                var result = await this.client.WaitForOrchestrationAsync(instance, DefaultTimeout);

                // Check that custom properties were extracted
                Assert.AreEqual(OrchestrationStatus.Failed, result.OrchestrationStatus);
                Assert.IsNotNull(result.FailureDetails);
                Assert.IsNotNull(result.FailureDetails.Properties);

                // Check the properties match the ArgumentOutOfRangeException.
                Assert.AreEqual("count", result.FailureDetails.Properties["Name"]);
                Assert.AreEqual("100", result.FailureDetails.Properties["Value"]);
            }
            finally
            {
                await this.worker.StopAsync();
            }
        }

        class ThrowCustomExceptionOrchestration : TaskOrchestration<string, string>
        {
            public override async Task<string> RunTask(OrchestrationContext context, string input)
            {
                await context.ScheduleTask<string>(typeof(ThrowCustomBusinessExceptionActivity), input);
                return "This should never be reached";
            }
        }

        class ThrowCustomBusinessExceptionActivity : TaskActivity<string, string>
        {
            protected override string Execute(TaskContext context, string input)
            {
                throw new CustomBusinessException("Payment processing failed", "user123", "OrderProcessing");
            }
        }

        class SimpleThrowExceptionOrchestration : TaskOrchestration<string, string>
        {
            public override Task<string> RunTask(OrchestrationContext context, string input)
            {
                throw new ArgumentOutOfRangeException("count", 100, "Count is not valid.");
            }
        }

        class SubOrchestrationThrowExceptionOrchestration : TaskOrchestration<string, string>
        {
            public override async Task<string> RunTask(OrchestrationContext context, string input)
            {
                await context.CreateSubOrchestrationInstance<string>(typeof(ThrowArgumentOutofRangeExceptionASubOrchestration), input);
                return "This should never be reached";
            }
        }

        class ThrowArgumentOutofRangeExceptionASubOrchestration : TaskOrchestration<string, string>
        {
            public override async Task<string> RunTask(OrchestrationContext context, string input)
            {
                await context.ScheduleTask<string>(typeof(ThrowArgumentOutofRangeExceptionActivity), input);
                return "This should never be reached";
            }
        }

        class ThrowArgumentOutofRangeExceptionActivity : TaskActivity<string, string>
        {
            protected override string Execute(TaskContext context, string input)
            {
                throw new ArgumentOutOfRangeException("count", 100, "Count is not valid.");
            }
        }

        class ThrowInvalidOperationExceptionOrchestration : TaskOrchestration<string, string>
        {
            public override async Task<string> RunTask(OrchestrationContext context, string input)
            {
                await context.ScheduleTask<string>(typeof(ThrowInvalidOperationExceptionActivity), input);
                return "This should never be reached";
            }
        }

        class ThrowInvalidOperationExceptionActivity : TaskActivity<string, string>
        {
            protected override string Execute(TaskContext context, string input)
            {
                throw new InvalidOperationException("This is a test exception");
            }
        }

        // Test exception with custom properties
        [Serializable]
        class CustomBusinessException : Exception
        {
            public string UserId { get; }
            public string BusinessContext { get; }
            public string? TestNullObject { get; }

            public CustomBusinessException(string message, string userId, string businessContext)
                : base(message)
            {
                UserId = userId;
                BusinessContext = businessContext;
                TestNullObject = null; // Explicitly set to null for testing
            }

            protected CustomBusinessException(SerializationInfo info, StreamingContext context)
                : base(info, context)
            {
                UserId = info.GetString(nameof(UserId)) ?? string.Empty;
                BusinessContext = info.GetString(nameof(BusinessContext)) ?? string.Empty;
                TestNullObject = info.GetString(nameof(TestNullObject)); // This will be null
            }

            public override void GetObjectData(SerializationInfo info, StreamingContext context)
            {
                base.GetObjectData(info, context);
                info.AddValue(nameof(UserId), UserId);
                info.AddValue(nameof(BusinessContext), BusinessContext);
                info.AddValue(nameof(TestNullObject), TestNullObject);
            }
        }

        // Test provider that includes null values in different ways
        class TestExceptionPropertiesProvider : IExceptionPropertiesProvider
        {
            public IDictionary<string, object?>? GetExceptionProperties(Exception exception)
            {
                return exception switch
                {
                    ArgumentOutOfRangeException e => new Dictionary<string, object?>
                    {
                        ["Name"] = e.ParamName ?? string.Empty,
                        ["Value"] = e.ActualValue?.ToString() ?? string.Empty,
                    },
                    CustomBusinessException businessEx => new Dictionary<string, object?>
                    {
                        ["ExceptionTypeName"] = nameof(CustomBusinessException),
                        ["UserId"] = businessEx.UserId,
                        ["BusinessContext"] = businessEx.BusinessContext,
                        ["Timestamp"] = DateTime.UtcNow,
                        ["TestNullObject"] = businessEx.TestNullObject, // This comes from the exception property (null)
                        ["DirectNullValue"] = null, // This is directly set to null
                        ["EmptyString"] = string.Empty // Non-null value for comparison
                    },
                    _ => null // No custom properties for other exceptions
                };
            }
        }

        [Serializable]
        class CustomException : Exception
        {
            public CustomException(string message)
                : base(message)
            {
            }

            protected CustomException(SerializationInfo info, StreamingContext context)
                : base(info, context)
            {
            }
        }

        [TestMethod]
        public void IsCausedBy_DoesNotThrow_WhenMultipleAssembliesDefineSameType()
        {
            // Create two dynamic assemblies, each containing an Exception-derived type with the
            // same fully qualified name. This simulates the scenario where the same exception type
            // is loaded from multiple assemblies (e.g. different NuGet package versions).
            string typeName = "TestDynamic.DuplicateException";
            CreateDynamicAssemblyWithExceptionType(typeName, "DynAssembly1");
            CreateDynamicAssemblyWithExceptionType(typeName, "DynAssembly2");

            // Create a FailureDetails whose ErrorType won't be resolved by Type.GetType(),
            // typeof(T).Assembly, or the calling assembly, forcing the AppDomain fallback path.
            var details = new FailureDetails(
                typeName, "Test error", stackTrace: null, innerFailure: null, isNonRetriable: false);

            // The old implementation would either throw AmbiguousMatchException or return false
            // when multiple assemblies contained the same type. The fix uses Any() so this should
            // succeed without throwing.
            bool result = details.IsCausedBy<Exception>();

            Assert.IsTrue(result);
        }

        static void CreateDynamicAssemblyWithExceptionType(string typeName, string assemblyName)
        {
            var asmName = new AssemblyName(assemblyName);
            var asmBuilder = AssemblyBuilder.DefineDynamicAssembly(asmName, AssemblyBuilderAccess.Run);
            var modBuilder = asmBuilder.DefineDynamicModule(assemblyName);
            var typeBuilder = modBuilder.DefineType(typeName, TypeAttributes.Public, typeof(Exception));
            typeBuilder.CreateType();
        }
    }
}