From 22f1ec23d6fc73e4a9bb1dd9df7dff75d8431427 Mon Sep 17 00:00:00 2001
From: pinzart90 <tiberiu.pinzariu@gmail.com>
Date: Wed, 1 Jan 2025 22:20:59 -0500
Subject: [PATCH] Update NotificationsExtensionTests.cs

---
 .../NotificationsExtensionTests.cs            | 24 +++++++++++--------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/test/DynamoCoreWpfTests/ViewExtensions/NotificationsExtensionTests.cs b/test/DynamoCoreWpfTests/ViewExtensions/NotificationsExtensionTests.cs
index 54915f00909..574ed58f54d 100644
--- a/test/DynamoCoreWpfTests/ViewExtensions/NotificationsExtensionTests.cs
+++ b/test/DynamoCoreWpfTests/ViewExtensions/NotificationsExtensionTests.cs
@@ -22,21 +22,25 @@ public void PressNotificationButtonAndShowPopup()
             notificationsButton.RaiseEvent(new RoutedEventArgs(ButtonBase.ClickEvent));
 
             var notificationExtension = this.View.viewExtensionManager.ViewExtensions.OfType<NotificationsViewExtension>().FirstOrDefault();
+            NotificationUI notificationUI = null;
+
             // Wait for the NotificationCenterController webview2 control to finish initialization
             DispatcherUtil.DoEventsLoop(() =>
             {
-                return notificationExtension.notificationCenterController.initState == DynamoUtilities.AsyncMethodState.Done;
+                if (notificationExtension.notificationCenterController.initState == DynamoUtilities.AsyncMethodState.Done)
+                {
+                    notificationUI = PresentationSource.CurrentSources.OfType<System.Windows.Interop.HwndSource>()
+                            .Select(h => h.RootVisual)
+                            .OfType<FrameworkElement>()
+                            .Select(f => f.Parent)
+                            .OfType<NotificationUI>()
+                            .FirstOrDefault(p => p.IsOpen);
+
+                    return notificationUI != null;
+                }
+                return false;
             });
             
-            Assert.AreEqual(DynamoUtilities.AsyncMethodState.Done, notificationExtension.notificationCenterController.initState);
-
-            NotificationUI notificationUI = PresentationSource.CurrentSources.OfType<System.Windows.Interop.HwndSource>()
-                                        .Select(h => h.RootVisual)
-                                        .OfType<FrameworkElement>()
-                                        .Select(f => f.Parent)
-                                        .OfType<NotificationUI>()
-                                        .FirstOrDefault(p => p.IsOpen);
-
             Assert.NotNull(notificationUI, "Notification popup not part of the dynamo visual tree");
             var webView = notificationUI.FindName("webView");
             Assert.NotNull(webView, "WebView framework element not found.");