From 1aa6421aa89d5ae2896e25cb26b130a16f6ceda9 Mon Sep 17 00:00:00 2001
From: Ivan Khvostishkov <87804596+ivan-khvostishkov@users.noreply.github.com>
Date: Wed, 23 Aug 2023 18:42:29 +0200
Subject: [PATCH] Better logic for IDE

---
 tests/js/drop_studio_file.js |  2 +-
 tests/test_ide.py            | 27 ++++++++++++---------------
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/tests/js/drop_studio_file.js b/tests/js/drop_studio_file.js
index f0b4e25..0fd3124 100644
--- a/tests/js/drop_studio_file.js
+++ b/tests/js/drop_studio_file.js
@@ -32,7 +32,7 @@ input.onchange = function () {
         }
     }
 
-    ['dragenter', 'drop'].forEach(function (name) {
+    ['dragenter', 'drop', 'dragleave'].forEach(function (name) {
         console.log('Dispatching mouse event with file transfer')
         const evt = document.createEvent('MouseEvent');
         evt.initMouseEvent(name, !0, !0, window, 0, 0, 0, x, y, !1, !1, !1, !1, 0, null);
diff --git a/tests/test_ide.py b/tests/test_ide.py
index 279b7b8..ccf04d7 100644
--- a/tests/test_ide.py
+++ b/tests/test_ide.py
@@ -333,9 +333,6 @@ def test_studio_notebook_in_firefox(request):
     logging.info("Launching SageMaker Studio")
     browser.get(studio_pre_signed_url)
 
-    logging.info("Checking for SageMaker Studio in Firefox")
-    assert "JupyterLab" in browser.title
-
     logging.info("Waiting for SageMaker Studio to launch")
     WebDriverWait(browser, 600).until(
         EC.presence_of_element_located((By.XPATH, "//div[@id='space-menu']"))
@@ -344,7 +341,16 @@ def test_studio_notebook_in_firefox(request):
     logging.info(f"Found SageMaker Studio space menu item: {kernel_menu_item.text}")
     assert kernel_menu_item.text == 'test-data-science / Personal Studio'
 
-    time.sleep(10)  # wait until obscurity of the menu items is gone and UI is fully loaded
+    time.sleep(60)  # wait until obscurity of the menu items is gone and UI is fully loaded
+
+    logging.info("Checking the kernel name")
+    kernel_item = browser.find_element(
+        By.XPATH,
+        "//button[@class='bp3-button bp3-minimal jp-Toolbar-kernelName "
+        "jp-ToolbarButtonComponent minimal jp-Button']"
+    )
+    logging.info(f"Found Kernel name: {kernel_item.text}")
+    assert kernel_item.text == "Data Science 2.0\n|\nPython 3\n|\n2 vCPU + 8 GiB"
 
     dist_file_name_pattern = 'sagemaker_ssh_helper-.*-py3-none-any.whl'
     dist_file_name = [f for f in os.listdir('../dist') if re.match(dist_file_name_pattern, f)][0]
@@ -377,15 +383,6 @@ def test_studio_notebook_in_firefox(request):
     logging.info(f"Found SageMaker Studio restart kernel menu item: {restart_menu_item.text}")
     restart_menu_item.click()
 
-    logging.info("Checking the kernel name")
-    kernel_item = browser.find_element(
-        By.XPATH,
-        "//button[@class='bp3-button bp3-minimal jp-Toolbar-kernelName "
-        "jp-ToolbarButtonComponent minimal jp-Button']"
-    )
-    logging.info(f"Found Kernel name: {kernel_item.text}")
-    assert kernel_item.text == "Data Science 2.0\n|\nPython 3\n|\n2 vCPU + 8 GiB"
-
     # TODO: check banner if kernel is still starting, wait until banner disappears, then click restart
     # <div class="css-a7sx0c-bannerContainer sagemaker-starting-banner" id="sagemaker-notebook-banner"><div class="css-1qyc1pu-kernelStartingBannerContainer"><div><div class="css-6wrpfe-bannerSpinDiv"></div></div><div><p class="css-g9mx5z-bannerPromptSpanTitle">Starting notebook kernel...</p></div></div></div>
 
@@ -437,10 +434,10 @@ def upload_file(browser, file_abs_path):
     logging.info(f"Created a file upload item: {file_input}")
     file_input.send_keys(file_abs_path)
     time.sleep(5)  # Give time to overwrite dialog to apper
-    confirmOverride(browser)
+    confirm_overwrite(browser)
 
 
-def confirmOverride(browser):
+def confirm_overwrite(browser):
     overwrite_button = browser.find_elements(
         By.XPATH,
         "//div[@class='jp-Dialog-buttonLabel' "