diff --git a/GFramework.Core.Tests/state/StateMachineSystemTests.cs b/GFramework.Core.Tests/state/StateMachineSystemTests.cs
index b05ddc5..cb9facc 100644
--- a/GFramework.Core.Tests/state/StateMachineSystemTests.cs
+++ b/GFramework.Core.Tests/state/StateMachineSystemTests.cs
@@ -111,8 +111,8 @@ public class StateMachineSystemTests
         _stateMachine!.Register(state1);
         _stateMachine.Register(state2);
 
-        Assert.That(state1.GetContext(), Is.Null);
-        Assert.That(state2.GetContext(), Is.Null);
+        Assert.Throws<InvalidOperationException>(() => state1.GetContext());
+        Assert.Throws<InvalidOperationException>(() => state2.GetContext());
 
         _stateMachine.Init();