From f5a8eb40d0d14fec1972e89a0e4bc7954cf4a512 Mon Sep 17 00:00:00 2001
From: Nikita Evdokimov <nikita-e57@yandex.ru>
Date: Thu, 25 Jan 2024 10:54:15 +0300
Subject: [PATCH] TECH: fix flaky test

---
 .../SystemDialogSafetyObjectDisabledSanityTest.kt | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/samples/kautomator-sample/src/androidTest/java/com/kaspersky/kaspresso/kautomatorsample/kaspresso/sanity/dialogs_safety/SystemDialogSafetyObjectDisabledSanityTest.kt b/samples/kautomator-sample/src/androidTest/java/com/kaspersky/kaspresso/kautomatorsample/kaspresso/sanity/dialogs_safety/SystemDialogSafetyObjectDisabledSanityTest.kt
index 6b91327f3..56f918966 100644
--- a/samples/kautomator-sample/src/androidTest/java/com/kaspersky/kaspresso/kautomatorsample/kaspresso/sanity/dialogs_safety/SystemDialogSafetyObjectDisabledSanityTest.kt
+++ b/samples/kautomator-sample/src/androidTest/java/com/kaspersky/kaspresso/kautomatorsample/kaspresso/sanity/dialogs_safety/SystemDialogSafetyObjectDisabledSanityTest.kt
@@ -3,6 +3,7 @@ package com.kaspersky.kaspresso.kautomatorsample.kaspresso.sanity.dialogs_safety
 import android.os.Build
 import androidx.test.ext.junit.rules.activityScenarioRule
 import androidx.test.uiautomator.StaleObjectException
+import com.kaspersky.components.kautomator.intercept.exception.UnfoundedUiObjectException
 import com.kaspersky.kaspresso.kaspresso.Kaspresso
 import com.kaspersky.kaspresso.kautomatorsample.screen.SystemDialogsScreen
 import com.kaspersky.kaspresso.kautomatorsample.systemdialogs.SystemDialogsActivity
@@ -43,13 +44,21 @@ class SystemDialogSafetyObjectDisabledSanityTest : TestCase(
             SystemDialogsScreen {
                 btn2 {
                     flakySafely(3000) {
-                        Assert.assertThrows(null, StaleObjectException::class.java) {
-                            isDisplayed()
-                            click()
+                        Assert.assertThrows(SampleException::class.java) {
+                            try {
+                                isDisplayed()
+                                click()
+                            } catch (ex: StaleObjectException) {
+                                throw SampleException()
+                            } catch (ex: UnfoundedUiObjectException) {
+                                throw SampleException()
+                            }
                         }
                     }
                 }
             }
         }
     }
+
+    class SampleException : Throwable()
 }