From 4cb59830b5d52ea7528a4332f3d112a440b88bf9 Mon Sep 17 00:00:00 2001
From: David <github.site@torick.net>
Date: Mon, 22 Nov 2021 00:09:42 -0500
Subject: [PATCH] fix: Fix support of the CustomVisualStateMamager

---
 .../Given_VisualStateManager.cs               | 28 ++++++++++++++++++-
 src/Uno.UI/UI/Xaml/VisualStateManager.cs      | 28 ++++++++-----------
 2 files changed, 39 insertions(+), 17 deletions(-)

diff --git a/src/Uno.UI.Tests/Windows_UI_Xaml/VisualStateManagerTests/Given_VisualStateManager.cs b/src/Uno.UI.Tests/Windows_UI_Xaml/VisualStateManagerTests/Given_VisualStateManager.cs
index 0b530278ec8c..c9c715397eb5 100644
--- a/src/Uno.UI.Tests/Windows_UI_Xaml/VisualStateManagerTests/Given_VisualStateManager.cs
+++ b/src/Uno.UI.Tests/Windows_UI_Xaml/VisualStateManagerTests/Given_VisualStateManager.cs
@@ -234,6 +234,20 @@ public void When_TransitionAndSetter_InCompatibilityMode_Then_SetterAppliedBefor
 			}
 		}
 
+		[TestMethod]
+		public void When_CustomManager_Then_UseIt()
+		{
+			var (control, group) = Setup();
+			var vsm = new CustomManager();
+			VisualStateManager.SetCustomVisualStateManager((FrameworkElement)control.TemplatedRoot, vsm);
+
+			VisualStateManager.GoToState(control, "state1", true);
+			VisualStateManager.GoToState(control, "state2", true);
+
+			Assert.IsTrue(new[] { "state1", "state2" }.SequenceEqual(vsm.States));
+		}
+
+
 		private static (Control control, VisualStateGroup states) Setup()
 		{
 			var control = new Control { Name = "control", Tag = "initial", Template = new ControlTemplate(() => new Grid()) };
@@ -247,7 +261,7 @@ private static (Control control, VisualStateGroup states) Setup()
 		}
 
 		private static VisualState State(int id)
-			=> new VisualState
+			=> new()
 			{
 				Name = "state" + id,
 				Setters = {new Setter(new TargetPropertyPath("control", "Tag"), "state" + id)}
@@ -305,5 +319,17 @@ internal void Unset()
 				IsActive = false;
 			}
 		}
+
+		private class CustomManager : VisualStateManager
+		{
+			public List<string> States { get; } = new();
+
+			/// <inheritdoc />
+			protected override bool GoToStateCore(Control control, FrameworkElement templateRoot, string stateName, VisualStateGroup @group, VisualState state, bool useTransitions)
+			{
+				States.Add(stateName);
+				return base.GoToStateCore(control, templateRoot, stateName, @group, state, useTransitions);
+			}
+		}
 	}
 }
diff --git a/src/Uno.UI/UI/Xaml/VisualStateManager.cs b/src/Uno.UI/UI/Xaml/VisualStateManager.cs
index 6cc92d2e78b2..c230c7009c2e 100644
--- a/src/Uno.UI/UI/Xaml/VisualStateManager.cs
+++ b/src/Uno.UI/UI/Xaml/VisualStateManager.cs
@@ -160,26 +160,22 @@ public static bool GoToState(Control control, string stateName, bool useTransiti
 				return false;
 			}
 
-			var vsm = GetCustomVisualStateManager(control) ?? GetVisualStateManager(control);
-			if (vsm is null)
-			{
-				if (_log.IsEnabled(Uno.Foundation.Logging.LogLevel.Debug))
-				{
-					_log.DebugFormat("Failed to set state [{0}], there is no VisualStateManager on [{1}]", stateName, control);
-				}
+#if __WASM__
+			TryAssignDOMVisualStates(groups, templateRoot);
+#endif
 
-				return false;
+			if (templateRoot is not FrameworkElement fwRoot)
+			{
+				// For backward compatibility!
+				return GetVisualStateManager(control).GoToStateCorePrivateBaseImplementation(control, group, state, useTransitions);
 			}
 
+			// Note: We resolve the 'CustomVisualStateManager' on the 'fwRoot' like UWP,
+			//		 but for compatibility reason we resolve the default visual state manager on the 'control' itself.
+			//		 We should validate the behavior on UWP, including for controls that does not have templates!
+			var vsm = GetCustomVisualStateManager(fwRoot) ?? GetVisualStateManager(control);
 
-			var output = templateRoot is FrameworkElement fwRoot
-				? vsm.GoToStateCore(control, fwRoot, stateName, group, state, useTransitions)
-				: vsm.GoToStateCorePrivateBaseImplementation(control, group, state, useTransitions); // For backward compatibility!
-				
-#if __WASM__
-			TryAssignDOMVisualStates(groups, templateRoot);
-#endif
-			return output;
+			return vsm.GoToStateCore(control, fwRoot, stateName, group, state, useTransitions);
 		}
 
 		protected virtual bool GoToStateCore(Control control, FrameworkElement templateRoot, string stateName, VisualStateGroup group, VisualState state, bool useTransitions)