From 06871f2f00c03b439e84e140ab8558cdae826e48 Mon Sep 17 00:00:00 2001
From: "Nikita.Ashihmin" <nikita.ashihmin@jetbrains.com>
Date: Wed, 11 Sep 2024 00:08:21 +0400
Subject: [PATCH] PY-72902 Jupyter(fix): fix tests

GitOrigin-RevId: 763ad97d4a759ed215c4232b9ee0881915056a46
---
 .../context/NotebookDataContext.kt              |  2 +-
 .../context/NotebookDataContextUtils.kt         | 17 ++++++++++-------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContext.kt b/notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContext.kt
index 3bc22d5a9470..ed6d072c363c 100644
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContext.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContext.kt
@@ -12,7 +12,7 @@ object NotebookDataContext {
 
   val DataContext.notebookEditor: EditorImpl?
     get() {
-      val component = getData(PlatformCoreDataKeys.CONTEXT_COMPONENT) ?: return null
+      val component = getData(PlatformCoreDataKeys.CONTEXT_COMPONENT)
       val editor = getData(PlatformCoreDataKeys.EDITOR)
       val noteEditor = NotebookDataContextUtils.getCurrentEditor(editor, component) ?: return null
       if (NotebookDataContextUtils.hasFocusedSearchReplaceComponent(noteEditor, component))
diff --git a/notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContextUtils.kt b/notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContextUtils.kt
index 0cf8568672b6..787a2f3e03e6 100644
--- a/notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContextUtils.kt
+++ b/notebooks/visualization/src/com/intellij/notebooks/visualization/context/NotebookDataContextUtils.kt
@@ -3,6 +3,7 @@ package com.intellij.notebooks.visualization.context
 import com.intellij.find.SearchReplaceComponent
 import com.intellij.ide.IdeEventQueue
 import com.intellij.notebooks.visualization.NotebookCellLines
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.ex.EditorGutterComponentEx
 import com.intellij.openapi.editor.impl.EditorComponentImpl
@@ -24,8 +25,14 @@ object NotebookDataContextUtils {
         return mouseEditorOffset
     }
 
+    if (contextComponent!=null) {
+      val calculatedLine = getRespectiveLineNumberInEditor(noteEditor, contextComponent)
+      if (calculatedLine!=null)
+        return calculatedLine
+    }
 
-
+    if (ApplicationManager.getApplication().isUnitTestMode)
+      return null
     val point = editor.contentComponent.mousePosition ?: return null
     val logicalPosition = editor.xyToLogicalPosition(point)
     return logicalPosition.line
@@ -52,11 +59,7 @@ object NotebookDataContextUtils {
     return noteEditor
   }
 
-  fun getCaretLine(editor: EditorImpl): Int {
-    return editor.caretModel.logicalPosition.line
-  }
-
-  fun getRespectiveLineNumberInEditor(editor: Editor, component: Component): Int? {
+  private fun getRespectiveLineNumberInEditor(editor: Editor, component: Component): Int? {
     val point = SwingUtilities.convertPoint(component, 0, component.height, editor.contentComponent)
     val documentLineCount = editor.document.lineCount
 
@@ -77,7 +80,7 @@ object NotebookDataContextUtils {
     return editor.logicalPositionToOffset(editor.xyToLogicalPosition(point))
   }
 
-  fun hasFocusedSearchReplaceComponent(editor: Editor, contextComponent: Component): Boolean {
+  fun hasFocusedSearchReplaceComponent(editor: Editor, contextComponent: Component?): Boolean {
     val searchReplaceComponent = editor.headerComponent.asSafely<SearchReplaceComponent>() ?: return false
     return contextComponent === searchReplaceComponent.searchTextComponent ||
            contextComponent === searchReplaceComponent.replaceTextComponent