#76 - Custom next state override

source/Lite.StateMachine.BenchmarkTests/BasicStateBenchmarks.cs

@@ -43,7 +43,7 @@ public async Task BasicStatesRunsAsync()
   }
 
   [Benchmark]
-  public void FlatStateMachineRunsSync()
+  public void BasicStatesRunsSync()
   {
     var maxCounter = CyclesBeforeExit;
     PropertyBag parameters = new()

source/Lite.StateMachine.Tests/StateTests/CustomStateTests.cs

@@ -0,0 +1,137 @@
+// Copyright Xeno Innovations, Inc. 2025
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Threading.Tasks;
+using Lite.StateMachine.Tests.TestData;
+using Lite.StateMachine.Tests.TestData.Services;
+using Lite.StateMachine.Tests.TestData.States.CustomStates;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+
+namespace Lite.StateMachine.Tests.StateTests;
+
+[TestClass]
+public class CustomStateTests : TestBase
+{
+  public TestContext TestContext { get; set; }
+
+  [TestMethod]
+  [DataRow(false, DisplayName = "Don't skip State3")]
+  [DataRow(true, DisplayName = "Skip State3")]
+  public async Task BasicState_Override_Executes_SuccessAsync(bool skipState3)
+  {
+    // Assemble with Dependency Injection
+    var services = new ServiceCollection()
+      //// Register Services
+      .AddLogging(InlineTraceLogger(LogLevel.None))
+      .AddSingleton<IMessageService, MessageService>()
+      .BuildServiceProvider();
+
+    var msgService = services.GetRequiredService<IMessageService>();
+    Func<Type, object?> factory = t => ActivatorUtilities.CreateInstance(services, t);
+
+    var ctxProperties = new PropertyBag()
+    {
+      { ParameterType.Counter, 0 },
+      { ParameterType.TestExitEarly, skipState3 },
+    };
+
+    var machine = new StateMachine<CustomStateId>(factory, null, isContextPersistent: true);
+    machine.RegisterState<State1>(CustomStateId.State1, CustomStateId.State2_Dummy);
+    machine.RegisterState<State2Dummy>(CustomStateId.State2_Dummy, CustomStateId.State3);
+    machine.RegisterState<State2Success>(CustomStateId.State2_Success, CustomStateId.State3);
+    machine.RegisterState<State3>(CustomStateId.State3);
+
+    // Act - Start your engine!
+    await machine.RunAsync(CustomStateId.State1, ctxProperties, cancellationToken: TestContext.CancellationToken);
+
+    // Assert Results
+    Assert.IsNotNull(machine);
+    Assert.IsNull(machine.Context);
+
+    Assert.AreEqual(1, msgService.Counter1);
+    Assert.AreEqual(0, msgService.Counter2, "State2Dummy should never enter");
+    Assert.AreEqual(skipState3 ? 0 : 1, msgService.Counter3);
+  }
+
+  [TestMethod]
+  public async Task BasicState_Overrides_ThrowUnregisteredException_Async()
+  {
+    // Assemble with Dependency Injection
+    var services = new ServiceCollection()
+      //// Register Services
+      .AddLogging(InlineTraceLogger(LogLevel.None))
+      .AddSingleton<IMessageService, MessageService>()
+      .BuildServiceProvider();
+
+    var msgService = services.GetRequiredService<IMessageService>();
+    Func<Type, object?> factory = t => ActivatorUtilities.CreateInstance(services, t);
+
+    var ctxProperties = new PropertyBag()
+    {
+      { ParameterType.Counter, 0 },
+      { ParameterType.TestUnregisteredTransition, true },
+    };
+
+    var machine = new StateMachine<CustomStateId>(factory, null, isContextPersistent: true);
+    machine.RegisterState<State1>(CustomStateId.State1, CustomStateId.State2_Dummy);
+    machine.RegisterState<State2Dummy>(CustomStateId.State2_Dummy, CustomStateId.State3);
+    machine.RegisterState<State2Success>(CustomStateId.State2_Success, CustomStateId.State3);
+    machine.RegisterState<State3>(CustomStateId.State3);
+
+    // Act - Start your engine!
+    await Assert.ThrowsExactlyAsync<UnregisteredStateTransitionException>(()
+      => machine.RunAsync(CustomStateId.State1, ctxProperties, null, TestContext.CancellationToken));
+
+    // Assert Results
+    Assert.IsNotNull(machine);
+    Assert.IsNull(machine.Context);
+
+    Assert.AreEqual(0, msgService.Counter1);
+    Assert.AreEqual(0, msgService.Counter2, "State2Dummy should never enter");
+  }
+
+  [TestMethod]
+  [DataRow(false, DisplayName = "Run State2_Sub3")]
+  [DataRow(true, DisplayName = "Skip State2_Sub3")]
+  public async Task Composite_Override_Executes_SuccessAsync(bool skipSubState3)
+  {
+    // Assemble with Dependency Injection
+    var services = new ServiceCollection()
+      //// Register Services
+      .AddLogging(InlineTraceLogger(LogLevel.Trace))
+      .AddSingleton<IMessageService, MessageService>()
+      .BuildServiceProvider();
+
+    var msgService = services.GetRequiredService<IMessageService>();
+    Func<Type, object?> factory = t => ActivatorUtilities.CreateInstance(services, t);
+
+    var ctxProperties = new PropertyBag()
+    {
+      { ParameterType.Counter, 0 },
+      { ParameterType.TestExitEarly2, skipSubState3 },
+    };
+
+    var machine = new StateMachine<CustomStateId>(factory, null, isContextPersistent: true);
+
+    machine.RegisterState<State1>(CustomStateId.State1, CustomStateId.State2_Dummy);
+    machine.RegisterState<State2Dummy>(CustomStateId.State2_Dummy, CustomStateId.State3);
+    machine.RegisterComposite<State2Success>(CustomStateId.State2_Success, CustomStateId.State2_Sub1, CustomStateId.State3);
+    machine.RegisterSubState<State2Success_Sub1>(CustomStateId.State2_Sub1, CustomStateId.State2_Success, CustomStateId.State2_Sub2);
+    machine.RegisterSubState<State2Success_Sub2>(CustomStateId.State2_Sub2, CustomStateId.State2_Success, CustomStateId.State2_Sub3);
+    machine.RegisterSubState<State2Success_Sub3>(CustomStateId.State2_Sub3, CustomStateId.State2_Success);
+    machine.RegisterState<State3>(CustomStateId.State3);
+
+    // Act - Start your engine!
+    await machine.RunAsync(CustomStateId.State1, ctxProperties, cancellationToken: TestContext.CancellationToken);
+
+    // Assert Results
+    Assert.IsNotNull(machine);
+    Assert.IsNull(machine.Context);
+
+    Assert.AreEqual(skipSubState3 ? 2 : 3, msgService.Counter1, "State Counter1 failed.");
+    Assert.AreEqual(0, msgService.Counter2, "State2Dummy should never enter");
+    Assert.AreEqual(1, msgService.Counter3, "Skip Substate Counter3 failed");
+  }
+}

source/Lite.StateMachine.Tests/DiTests/MsDiTests.cs → source/Lite.StateMachine.Tests/StateTests/DiMsTests.cs

@@ -13,11 +13,12 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Console;
 
-namespace Lite.StateMachine.Tests.DiTests;
+namespace Lite.StateMachine.Tests.StateTests;
 
+/// <summary>Microsoft Dependency Injection Tests.</summary>
 [TestClass]
 [System.Diagnostics.CodeAnalysis.SuppressMessage("StyleCop.CSharp.ReadabilityRules", "SA1124:Do not use regions", Justification = "Allowed for this test class")]
-public class MsDiTests
+public class DiMsTests
 {
   [TestMethod]
   public async Task Basic_FlatStates_SuccessTestAsync()
@@ -197,7 +198,7 @@ public async Task RegisterState_MsDi_EventAggregatorOnly_SuccessTestAsync()
 
     var aggregator = services.GetRequiredService<IEventAggregator>();
     var msgService = services.GetRequiredService<IMessageService>();
-    var logService = services.GetRequiredService<ILogger<MsDiTests>>();
+    var logService = services.GetRequiredService<ILogger<DiMsTests>>();
 
     msgService.Counter1 = 0;
     msgService.Counter2 = 0;

source/Lite.StateMachine.Tests/TestData/ParameterType.cs

@@ -3,9 +3,23 @@
 
 namespace Lite.StateMachine.Tests.TestData;
 
-public static class ParameterType
+public enum ParameterType
 {
-  public const string Counter = "Counter";
-  public const string HungStateAvoidance = "DoNotAllowHungStatesTest";
-  public const string KeyTest = "TestKey";
+  /// <summary>Generic counter.</summary>
+  Counter,
+
+  /// <summary>Tests for DoNotAllowHungStatesTest.</summary>
+  HungStateAvoidance,
+
+  /// <summary>Random test.</summary>
+  KeyTest,
+
+  /// <summary>Test triggers an early exit. Setting OnSuccess to NULL.</summary>
+  TestExitEarly,
+
+  /// <summary>Test triggers a 2nd early exit. Setting OnSuccess to NULL.</summary>
+  TestExitEarly2,
+
+  /// <summary>Test trigger to go to an invalid state transition.</summary>
+  TestUnregisteredTransition,
 }

source/Lite.StateMachine.Tests/TestData/States/StateEnums.cs → source/Lite.StateMachine.Tests/TestData/StateEnums.cs

@@ -1,7 +1,7 @@
 // Copyright Xeno Innovations, Inc. 2025
 // See the LICENSE file in the project root for more information.
 
-namespace Lite.StateMachine.Tests.TestData.States;
+namespace Lite.StateMachine.Tests.TestData;
 
 #pragma warning disable SA1649 // File name should match first type name
 
@@ -46,4 +46,24 @@ public enum CompositeMsgStateId
   Failure,
 }
 
+public enum CustomStateId
+{
+  /// <summary>Initial state.</summary>
+  State1,
+
+  /// <summary>Dummy state that is registered but never transitioned to.</summary>
+  State2_Dummy,
+
+  State2_Success,
+
+  /// <summary>This state must NEVER be registered.</summary>
+  State2_Unregistered,
+
+  State2_Sub1,
+  State2_Sub2,
+  State2_Sub3,
+
+  State3,
+}
+
 #pragma warning restore SA1649 // File name should match first type name

source/Lite.StateMachine.Tests/TestData/States/CustomBasicStates.cs

@@ -0,0 +1,139 @@
+// Copyright Xeno Innovations, Inc. 2025
+// See the LICENSE file in the project root for more information.
+
+using System.Threading.Tasks;
+using Lite.StateMachine.Tests.TestData.Services;
+
+#pragma warning disable SA1649 // File name should match first type name
+#pragma warning disable SA1402 // File may only contain a single type
+#pragma warning disable IDE0130 // Namespace does not match folder structure
+
+namespace Lite.StateMachine.Tests.TestData.States.CustomStates;
+
+public class State1 : StateBase<State1, CustomStateId>
+{
+  public State1() => HasDebugLogging = true;
+
+  public override async Task OnEnter(Context<CustomStateId> ctx)
+  {
+    int cnt = ctx.ParameterAsInt(ParameterType.Counter);
+
+    if (ctx.ParameterAsBool(ParameterType.TestUnregisteredTransition))
+      ctx.NextStates.OnSuccess = CustomStateId.State2_Unregistered;
+    else
+      ctx.NextStates.OnSuccess = CustomStateId.State2_Success;
+
+    await base.OnEnter(ctx);
+  }
+}
+
+/// <summary>This state should NEVER be transitioned into.</summary>
+public class State2Dummy : StateBase<State2Dummy, CustomStateId>
+{
+  private readonly IMessageService _msgService;
+
+  public State2Dummy(IMessageService msg)
+  {
+    _msgService = msg;
+    HasDebugLogging = true;
+  }
+
+  public override Task OnEnter(Context<CustomStateId> context)
+  {
+    Assert.Fail("Overridden state transitions should not all us to be here.");
+    _msgService.Counter2++;
+    return base.OnEnter(context);
+  }
+}
+
+public class State2Success : StateBase<State2Success, CustomStateId>
+{
+  private readonly IMessageService _msgService;
+
+  public State2Success(IMessageService msg)
+  {
+    _msgService = msg;
+    HasDebugLogging = true;
+  }
+
+  public override Task OnEnter(Context<CustomStateId> ctx)
+  {
+    _msgService.Counter1++;
+
+    if (ctx.ParameterAsBool(ParameterType.TestExitEarly))
+      ctx.NextStates.OnSuccess = null;
+
+    return base.OnEnter(ctx);
+  }
+
+  public override Task OnExit(Context<CustomStateId> context)
+  {
+    // When operating as a Composite, we need to pass back Success.
+    context.NextState(Result.Success);
+    return base.OnExit(context);
+  }
+}
+
+public class State2Success_Sub1 : StateBase<State2Success_Sub1, CustomStateId>
+{
+  public State2Success_Sub1() => HasDebugLogging = true;
+}
+
+public class State2Success_Sub2 : StateBase<State2Success_Sub2, CustomStateId>
+{
+  private readonly IMessageService _msgService;
+
+  public State2Success_Sub2(IMessageService msg)
+  {
+    _msgService = msg;
+    HasDebugLogging = true;
+  }
+
+  public override Task OnEnter(Context<CustomStateId> ctx)
+  {
+    _msgService.Counter1++;
+
+    if (ctx.ParameterAsBool(ParameterType.TestExitEarly2))
+      ctx.NextStates.OnSuccess = null;
+
+    return base.OnEnter(ctx);
+  }
+}
+
+public class State2Success_Sub3 : StateBase<State2Success_Sub3, CustomStateId>
+{
+  private readonly IMessageService _msgService;
+
+  public State2Success_Sub3(IMessageService msg)
+  {
+    _msgService = msg;
+    HasDebugLogging = true;
+  }
+
+  public override Task OnEnter(Context<CustomStateId> ctx)
+  {
+    _msgService.Counter1++;
+    return base.OnEnter(ctx);
+  }
+}
+
+public class State3 : StateBase<State3, CustomStateId>
+{
+  private readonly IMessageService _msgService;
+
+  public State3(IMessageService msg)
+  {
+    _msgService = msg;
+    HasDebugLogging = true;
+  }
+
+  public override Task OnEnter(Context<CustomStateId> ctx)
+  {
+    _msgService.Counter3++;
+    return base.OnEnter(ctx);
+  }
+}
+
+#pragma warning restore IDE0130 // Namespace does not match folder structure
+#pragma warning restore SA1649 // File name should match first type name
+#pragma warning restore SA1402 // File may only contain a single type