From 408460a7e88ebb585b9e7e0088a41044150d253f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miko=C5=82aj=20Milewski?= <mikolaj.milewski@gmail.com>
Date: Tue, 27 Aug 2024 11:31:50 +0200
Subject: [PATCH] IHostApplicationLifetime for testing

---
 .../Stateflows.Testing/StateflowsTestClass.cs | 10 +++++++
 Core/Stateflows.Testing/TestingHost.cs        | 28 +++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 Core/Stateflows.Testing/TestingHost.cs

diff --git a/Core/Stateflows.Testing/StateflowsTestClass.cs b/Core/Stateflows.Testing/StateflowsTestClass.cs
index 71f7726..0f0c71d 100644
--- a/Core/Stateflows.Testing/StateflowsTestClass.cs
+++ b/Core/Stateflows.Testing/StateflowsTestClass.cs
@@ -12,6 +12,7 @@
 using Stateflows.Common.Registration.Interfaces;
 using Stateflows.Testing.StateMachines.Sequence;
 using Stateflows.Common.Classes;
+using Stateflows.Testing;
 
 namespace StateMachine.IntegrationTests.Utils
 {
@@ -29,15 +30,22 @@ public abstract class StateflowsTestClass
 
         protected ExecutionSequenceObserver ExecutionSequence => ServiceProvider.GetRequiredService<ExecutionSequenceObserver>();
 
+        private TestingHost testingHost;
+        private TestingHost TestingHost => testingHost ??= ServiceProvider.GetRequiredService<TestingHost>();
+
         public virtual void Initialize()
         {
             ServiceCollection.AddStateflows(b => InitializeStateflows(b));
             ServiceCollection.AddSingleton<IExecutionSequenceBuilder, ExecutionSequence>();
+            ServiceCollection.AddSingleton<TestingHost>();
+            ServiceCollection.AddSingleton<IHostApplicationLifetime>(services => services.GetRequiredService<TestingHost>());
             ServiceCollection.AddLogging(builder => builder.AddConsole());
 
             var hostedServices = ServiceProvider.GetRequiredService<IEnumerable<IHostedService>>();
             Task.WaitAll(hostedServices.Select(s => s.StartAsync(new CancellationToken())).ToArray());
 
+            TestingHost.StartApplication();
+
             ContextValues.GlobalValues.Clear();
             ContextValues.StateValues.Clear();
             ContextValues.SourceStateValues.Clear();
@@ -46,6 +54,8 @@ public virtual void Initialize()
 
         public virtual void Cleanup()
         {
+            TestingHost.StopApplication();
+
             _serviceCollection = null;
             _serviceProvider = null;
         }
diff --git a/Core/Stateflows.Testing/TestingHost.cs b/Core/Stateflows.Testing/TestingHost.cs
new file mode 100644
index 0000000..dd2d7e9
--- /dev/null
+++ b/Core/Stateflows.Testing/TestingHost.cs
@@ -0,0 +1,28 @@
+using System.Threading;
+using Microsoft.Extensions.Hosting;
+
+namespace Stateflows.Testing
+{
+    internal class TestingHost : IHostApplicationLifetime
+    {
+        private readonly CancellationTokenSource ApplicationStartedToken = new CancellationTokenSource();
+        public CancellationToken ApplicationStarted => ApplicationStartedToken.Token;
+
+        private readonly CancellationTokenSource ApplicationStoppingToken = new CancellationTokenSource();
+        public CancellationToken ApplicationStopping => ApplicationStoppingToken.Token;
+
+        private readonly CancellationTokenSource ApplicationStoppedToken = new CancellationTokenSource();
+        public CancellationToken ApplicationStopped => ApplicationStoppedToken.Token;
+
+        public void StartApplication()
+        {
+            ApplicationStartedToken.Cancel();
+        }
+
+        public void StopApplication()
+        {
+            ApplicationStoppingToken.Cancel();
+            ApplicationStoppedToken.Cancel();
+        }
+    }
+}