Update NotificationsExtensionTests.cs

DynamoDS · Jan 2, 2025 · 22f1ec2 · 22f1ec2
1 parent 213a64c
commit 22f1ec2
Showing 1 changed file with 14 additions and 10 deletions.
diff --git 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.");