Try to fix previous commit

Author: mcbarton

scripts/automated-notebook-run-script.py

@@ -88,6 +88,9 @@ def run_notebook(driver, notebook_area, args):
 
         if cell_is_waiting_for_input(driver):
             print("Cell requesting input")
+            current_cell = driver.find_element(
+                By.CSS_SELECTOR, ".jp-Notebook-cell.jp-mod-selected"
+            )
             input_box = WebDriverWait(driver, 5).until(
                 EC.visibility_of_element_located((By.CSS_SELECTOR, ".jp-Stdin-input"))
             )
@@ -125,11 +128,11 @@ def run_notebook(driver, notebook_area, args):
         time.sleep(0.1)
         if not cell_is_waiting_for_input(driver):
             wait_for_idle_status(driver, current_cell, start_time, args.timeout)
-        next_cell = current_cell.find_element(
-            By.XPATH,
-            "following-sibling::div[contains(@class,'jp-Notebook-cell')][1]",
-        )
-        current_cell = next_cell
+            next_cell = current_cell.find_element(
+                By.XPATH,
+                "following-sibling::div[contains(@class,'jp-Notebook-cell')][1]",
+            )
+            current_cell = next_cell
 
 def download_notebook(driver):
     """This function is used to download the notebook currently open."""