Repository: python-lsp/python-lsp-server
Pull request #466 (2023-10-17 23:09:24)

Ignore notebook names on cell completion for autoimport
## What is changed?

With the notebook document, every cell is its own document living in the workspace. When doing rope autoimport completions on a cell document position, we should not provide completion suggestions for names already imported before that (see video below).

## How is this tested?

- [x] integration test
- [x] manually on Databricks notebooks

https://github.com/python-lsp/python-lsp-server/assets/91616041/08965bd6-9921-41e2-b099-08e9d8244096



tkrabel-db: Thanks for the review @ccordoba12 . I added all your suggestions! 🚀 
tkrabel-db: Weird, `Static code analysis / Static code analysis (pull_request)` failed but when looking at the details, everything passed ...
krassowski: Kicked, let's see if it happens again.

Patch:
diff --git a/pylsp/hookspecs.py b/pylsp/hookspecs.py
index d732b1d..9c9cf38 100644
--- a/pylsp/hookspecs.py
+++ b/pylsp/hookspecs.py
@@ -25,7 +25,7 @@ def pylsp_commands(config, workspace):
 
 
 @hookspec
-def pylsp_completions(config, workspace, document, position):
+def pylsp_completions(config, workspace, document, position, ignored_names):
     pass
 
 
diff --git a/pylsp/plugins/rope_autoimport.py b/pylsp/plugins/rope_autoimport.py
index 6c4784a..1caab35 100644
--- a/pylsp/plugins/rope_autoimport.py
+++ b/pylsp/plugins/rope_autoimport.py
@@ -1,7 +1,7 @@
 # Copyright 2022- Python Language Server Contributors.
 
 import logging
-from typing import Any, Dict, Generator, List, Optional, Set
+from typing import Any, Dict, Generator, List, Optional, Set, Union
 
 import parso
 from jedi import Script
@@ -153,7 +153,11 @@ def get_names(script: Script) -> Set[str]:
 
 @hookimpl
 def pylsp_completions(
-    config: Config, workspace: Workspace, document: Document, position
+    config: Config,
+    workspace: Workspace,
+    document: Document,
+    position,
+    ignored_names: Union[Set[str], None],
 ):
     """Get autoimport suggestions."""
     line = document.lines[position["line"]]
@@ -164,7 +168,9 @@ def pylsp_completions(
     word = word_node.value
     log.debug(f"autoimport: searching for word: {word}")
     rope_config = config.settings(document_path=document.path).get("rope", {})
-    ignored_names: Set[str] = get_names(document.jedi_script(use_document_path=True))
+    ignored_names: Set[str] = ignored_names or get_names(
+        document.jedi_script(use_document_path=True)
+    )
     autoimport = workspace._rope_autoimport(rope_config)
     suggestions = list(autoimport.search_full(word, ignored_names=ignored_names))
     results = list(
diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
index 4c3ea0d..760ad97 100644
--- a/pylsp/python_lsp.py
+++ b/pylsp/python_lsp.py
@@ -394,7 +394,16 @@ class PythonLSPServer(MethodDispatcher):
         return flatten(self._hook("pylsp_code_lens", doc_uri))
 
     def completions(self, doc_uri, position):
-        completions = self._hook("pylsp_completions", doc_uri, position=position)
+        workspace = self._match_uri_to_workspace(doc_uri)
+        document = workspace.get_document(doc_uri)
+        ignored_names = None
+        if isinstance(document, Cell):
+            # We need to get the ignored names from the whole notebook document
+            notebook_document = workspace.get_maybe_document(document.notebook_uri)
+            ignored_names = notebook_document.jedi_names(doc_uri)
+        completions = self._hook(
+            "pylsp_completions", doc_uri, position=position, ignored_names=ignored_names
+        )
         return {"isIncomplete": False, "items": flatten(completions)}
 
     def completion_item_resolve(self, completion_item):
diff --git a/pylsp/workspace.py b/pylsp/workspace.py
index e0e16ef..fb524c7 100644
--- a/pylsp/workspace.py
+++ b/pylsp/workspace.py
@@ -595,6 +595,7 @@ class Notebook:
         self.version = version
         self.cells = cells or []
         self.metadata = metadata or {}
+        self._lock = RLock()
 
     def __str__(self):
         return "Notebook with URI '%s'" % str(self.uri)
@@ -625,6 +626,31 @@ class Notebook:
             offset += num_lines
         return cell_data
 
+    @lock
+    def jedi_names(
+        self,
+        up_to_cell_uri: Optional[str] = None,
+        all_scopes=False,
+        definitions=True,
+        references=False,
+    ):
+        """
+        Get the names in the notebook up to a certain cell.
+
+        Parameters
+        ----------
+        up_to_cell_uri: str, optional
+            The cell uri to stop at. If None, all cells are considered.
+        """
+        names = set()
+        for cell in self.cells:
+            cell_uri = cell["document"]
+            cell_document = self.workspace.get_cell_document(cell_uri)
+            names.update(cell_document.jedi_names(all_scopes, definitions, references))
+            if cell_uri == up_to_cell_uri:
+                break
+        return set(name.name for name in names)
+
 
 class Cell(Document):
     """


Test patch:
diff --git a/test/fixtures.py b/test/fixtures.py
index 03d0f82..ed6206a 100644
--- a/test/fixtures.py
+++ b/test/fixtures.py
@@ -5,9 +5,11 @@ import os
 from io import StringIO
 from unittest.mock import MagicMock
 
-from test.test_utils import ClientServerPair
+from test.test_utils import ClientServerPair, CALL_TIMEOUT_IN_SECONDS
 
 import pytest
+import pylsp_jsonrpc
+
 from pylsp_jsonrpc.dispatchers import MethodDispatcher
 from pylsp_jsonrpc.endpoint import Endpoint
 from pylsp_jsonrpc.exceptions import JsonRpcException
@@ -24,7 +26,6 @@ DOC = """import sys
 def main():
     print sys.stdin.read()
 """
-CALL_TIMEOUT_IN_SECONDS = 30
 
 
 class FakeEditorMethodsMixin:
@@ -175,8 +176,13 @@ def client_server_pair():
 
     yield (client_server_pair_obj.client, client_server_pair_obj.server)
 
-    shutdown_response = client_server_pair_obj.client._endpoint.request(
-        "shutdown"
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
-    assert shutdown_response is None
-    client_server_pair_obj.client._endpoint.notify("exit")
+    try:
+        shutdown_response = client_server_pair_obj.client._endpoint.request(
+            "shutdown"
+        ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+        assert shutdown_response is None
+        client_server_pair_obj.client._endpoint.notify("exit")
+    except pylsp_jsonrpc.exceptions.JsonRpcInvalidParams:
+        # SQLite objects created in a thread can only be used in that same thread.
+        # This exeception is raised when testing rope autoimport.
+        client_server_pair_obj.client._endpoint.notify("exit")
diff --git a/test/plugins/test_autoimport.py b/test/plugins/test_autoimport.py
index dbb6f7a..b1c4677 100644
--- a/test/plugins/test_autoimport.py
+++ b/test/plugins/test_autoimport.py
@@ -1,13 +1,16 @@
 # Copyright 2022- Python Language Server Contributors.
 
-from typing import Dict, List
-from unittest.mock import Mock
+from typing import Any, Dict, List
+from unittest.mock import Mock, patch
+
+from test.test_notebook_document import wait_for_condition
+from test.test_utils import send_initialize_request, send_notebook_did_open
 
 import jedi
 import parso
 import pytest
 
-from pylsp import lsp, uris
+from pylsp import IS_WIN, lsp, uris
 from pylsp.config.config import Config
 from pylsp.plugins.rope_autoimport import _get_score, _should_insert, get_names
 from pylsp.plugins.rope_autoimport import (
@@ -16,9 +19,17 @@ from pylsp.plugins.rope_autoimport import (
 from pylsp.plugins.rope_autoimport import pylsp_initialize
 from pylsp.workspace import Workspace
 
+
 DOC_URI = uris.from_fs_path(__file__)
 
 
+def contains_autoimport(suggestion: Dict[str, Any], module: str) -> bool:
+    """Checks if `suggestion` contains an autoimport for `module`."""
+    return suggestion.get("label", "") == module and "import" in suggestion.get(
+        "detail", ""
+    )
+
+
 @pytest.fixture(scope="session")
 def autoimport_workspace(tmp_path_factory) -> Workspace:
     "Special autoimport workspace. Persists across sessions to make in-memory sqlite3 database fast."
@@ -39,7 +50,9 @@ def completions(config: Config, autoimport_workspace: Workspace, request):
     com_position = {"line": 0, "character": position}
     autoimport_workspace.put_document(DOC_URI, source=document)
     doc = autoimport_workspace.get_document(DOC_URI)
-    yield pylsp_autoimport_completions(config, autoimport_workspace, doc, com_position)
+    yield pylsp_autoimport_completions(
+        config, autoimport_workspace, doc, com_position, None
+    )
     autoimport_workspace.rm_document(DOC_URI)
 
 
@@ -141,45 +154,13 @@ def test_autoimport_defined_name(config, workspace):
     com_position = {"line": 1, "character": 3}
     workspace.put_document(DOC_URI, source=document)
     doc = workspace.get_document(DOC_URI)
-    completions = pylsp_autoimport_completions(config, workspace, doc, com_position)
+    completions = pylsp_autoimport_completions(
+        config, workspace, doc, com_position, None
+    )
     workspace.rm_document(DOC_URI)
     assert not check_dict({"label": "List"}, completions)
 
 
-# This test has several large issues.
-# 1. autoimport relies on its sources being written to disk. This makes testing harder
-# 2. the hook doesn't handle removed files
-# 3. The testing framework cannot access the actual autoimport object so it cannot clear the cache
-# def test_autoimport_update_module(config: Config, workspace: Workspace):
-#     document2 = "SomethingYouShouldntWrite = 1"
-#     document = """SomethingYouShouldntWrit"""
-#     com_position = {
-#         "line": 0,
-#         "character": 3,
-#     }
-#     doc2_path = workspace.root_path + "/test_file_no_one_should_write_to.py"
-#     if os.path.exists(doc2_path):
-#         os.remove(doc2_path)
-#     DOC2_URI = uris.from_fs_path(doc2_path)
-#     workspace.put_document(DOC_URI, source=document)
-#     doc = workspace.get_document(DOC_URI)
-#     completions = pylsp_autoimport_completions(config, workspace, doc, com_position)
-#     assert len(completions) == 0
-#     with open(doc2_path, "w") as f:
-#         f.write(document2)
-#     workspace.put_document(DOC2_URI, source=document2)
-#     doc2 = workspace.get_document(DOC2_URI)
-#     pylsp_document_did_save(config, workspace, doc2)
-#     assert check_dict({"label": "SomethingYouShouldntWrite"}, completions)
-#     workspace.put_document(DOC2_URI, source="\n")
-#     doc2 = workspace.get_document(DOC2_URI)
-#     os.remove(doc2_path)
-#     pylsp_document_did_save(config, workspace, doc2)
-#     completions = pylsp_autoimport_completions(config, workspace, doc, com_position)
-#     assert len(completions) == 0
-#     workspace.rm_document(DOC_URI)
-
-
 class TestShouldInsert:
     def test_dot(self):
         assert not should_insert("""str.""", 4)
@@ -233,3 +214,74 @@ def test_get_names():
     """
     results = get_names(jedi.Script(code=source))
     assert results == set(["blah", "bleh", "e", "hello", "someone", "sfa", "a", "b"])
+
+
+@pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
+def test_autoimport_for_notebook_document(
+    client_server_pair,
+):
+    client, server = client_server_pair
+    send_initialize_request(client)
+
+    with patch.object(server._endpoint, "notify") as mock_notify:
+        # Expectations:
+        # 1. We receive an autoimport suggestion for "os" in the first cell because
+        #    os is imported after that.
+        # 2. We don't receive an autoimport suggestion for "os" in the second cell because it's
+        #    already imported in the second cell.
+        # 3. We don't receive an autoimport suggestion for "os" in the third cell because it's
+        #    already imported in the second cell.
+        # 4. We receive an autoimport suggestion for "sys" because it's not already imported
+        send_notebook_did_open(client, ["os", "import os\nos", "os", "sys"])
+        wait_for_condition(lambda: mock_notify.call_count >= 3)
+
+    server.m_workspace__did_change_configuration(
+        settings={
+            "pylsp": {"plugins": {"rope_autoimport": {"enabled": True, "memory": True}}}
+        }
+    )
+    rope_autoimport_settings = server.workspace._config.plugin_settings(
+        "rope_autoimport"
+    )
+    assert rope_autoimport_settings.get("enabled", False) is True
+    assert rope_autoimport_settings.get("memory", False) is True
+
+    # 1.
+    suggestions = server.completions("cell_1_uri", {"line": 0, "character": 2}).get(
+        "items"
+    )
+    assert any(
+        suggestion
+        for suggestion in suggestions
+        if contains_autoimport(suggestion, "os")
+    )
+
+    # 2.
+    suggestions = server.completions("cell_2_uri", {"line": 1, "character": 2}).get(
+        "items"
+    )
+    assert not any(
+        suggestion
+        for suggestion in suggestions
+        if contains_autoimport(suggestion, "os")
+    )
+
+    # 3.
+    suggestions = server.completions("cell_3_uri", {"line": 0, "character": 2}).get(
+        "items"
+    )
+    assert not any(
+        suggestion
+        for suggestion in suggestions
+        if contains_autoimport(suggestion, "os")
+    )
+
+    # 4.
+    suggestions = server.completions("cell_4_uri", {"line": 0, "character": 3}).get(
+        "items"
+    )
+    assert any(
+        suggestion
+        for suggestion in suggestions
+        if contains_autoimport(suggestion, "sys")
+    )
diff --git a/test/test_language_server.py b/test/test_language_server.py
index 280a62f..401b1ce 100644
--- a/test/test_language_server.py
+++ b/test/test_language_server.py
@@ -5,7 +5,7 @@ import os
 import time
 import sys
 
-from test.test_utils import ClientServerPair
+from test.test_utils import ClientServerPair, send_initialize_request
 
 from flaky import flaky
 from pylsp_jsonrpc.exceptions import JsonRpcMethodNotFound
@@ -73,14 +73,7 @@ def test_not_exit_without_check_parent_process_flag(
     client_server_pair,
 ):
     client, _ = client_server_pair
-    response = client._endpoint.request(
-        "initialize",
-        {
-            "processId": 1234,
-            "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
-        },
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    response = send_initialize_request(client)
     assert "capabilities" in response
 
 
diff --git a/test/test_notebook_document.py b/test/test_notebook_document.py
index 15f187d..c63d279 100644
--- a/test/test_notebook_document.py
+++ b/test/test_notebook_document.py
@@ -1,10 +1,14 @@
 # Copyright 2021- Python Language Server Contributors.
 
-import os
 import time
-
 from unittest.mock import patch, call
-from test.fixtures import CALL_TIMEOUT_IN_SECONDS
+
+from test.test_utils import (
+    CALL_TIMEOUT_IN_SECONDS,
+    send_initialize_request,
+    send_notebook_did_open,
+)
+
 import pytest
 from pylsp.workspace import Notebook
 
@@ -24,14 +28,7 @@ def wait_for_condition(condition, timeout=CALL_TIMEOUT_IN_SECONDS):
 @pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
 def test_initialize(client_server_pair):
     client, server = client_server_pair
-    response = client._endpoint.request(
-        "initialize",
-        {
-            "processId": 1234,
-            "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
-        },
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    response = send_initialize_request(client)
     assert server.workspace is not None
     selector = response["capabilities"]["notebookDocumentSync"]["notebookSelector"]
     assert isinstance(selector, list)
@@ -41,13 +38,7 @@ def test_initialize(client_server_pair):
 def test_workspace_did_change_configuration(client_server_pair):
     """Test that we can update a workspace config w/o error when a notebook is open."""
     client, server = client_server_pair
-    client._endpoint.request(
-        "initialize",
-        {
-            "processId": 1234,
-            "rootPath": os.path.dirname(__file__),
-        },
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    send_initialize_request(client)
     assert server.workspace is not None
 
     with patch.object(server._endpoint, "notify") as mock_notify:
@@ -95,74 +86,12 @@ def test_notebook_document__did_open(
     client_server_pair,
 ):
     client, server = client_server_pair
-    client._endpoint.request(
-        "initialize",
-        {
-            "processId": 1234,
-            "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
-        },
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    send_initialize_request(client)
 
     with patch.object(server._endpoint, "notify") as mock_notify:
-        client._endpoint.notify(
-            "notebookDocument/didOpen",
-            {
-                "notebookDocument": {
-                    "uri": "notebook_uri",
-                    "notebookType": "jupyter-notebook",
-                    "cells": [
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_1_uri",
-                        },
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_2_uri",
-                        },
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_3_uri",
-                        },
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_4_uri",
-                        },
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_5_uri",
-                        },
-                    ],
-                },
-                # Test as many edge cases as possible for the diagnostics message
-                "cellTextDocuments": [
-                    {
-                        "uri": "cell_1_uri",
-                        "languageId": "python",
-                        "text": "",
-                    },
-                    {
-                        "uri": "cell_2_uri",
-                        "languageId": "python",
-                        "text": "\n",
-                    },
-                    {
-                        "uri": "cell_3_uri",
-                        "languageId": "python",
-                        "text": "\nimport sys\n\nabc\n\n",
-                    },
-                    {
-                        "uri": "cell_4_uri",
-                        "languageId": "python",
-                        "text": "x",
-                    },
-                    {
-                        "uri": "cell_5_uri",
-                        "languageId": "python",
-                        "text": "y\n",
-                    },
-                ],
-            },
+        # Test as many edge cases as possible for the diagnostics messages
+        send_notebook_did_open(
+            client, ["", "\n", "\nimport sys\n\nabc\n\n", "x", "y\n"]
         )
         wait_for_condition(lambda: mock_notify.call_count >= 5)
         expected_call_args = [
@@ -259,48 +188,11 @@ def test_notebook_document__did_change(
     client_server_pair,
 ):
     client, server = client_server_pair
-    client._endpoint.request(
-        "initialize",
-        {
-            "processId": 1234,
-            "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
-        },
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    send_initialize_request(client)
 
     # Open notebook
     with patch.object(server._endpoint, "notify") as mock_notify:
-        client._endpoint.notify(
-            "notebookDocument/didOpen",
-            {
-                "notebookDocument": {
-                    "uri": "notebook_uri",
-                    "notebookType": "jupyter-notebook",
-                    "cells": [
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_1_uri",
-                        },
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_2_uri",
-                        },
-                    ],
-                },
-                "cellTextDocuments": [
-                    {
-                        "uri": "cell_1_uri",
-                        "languageId": "python",
-                        "text": "import sys",
-                    },
-                    {
-                        "uri": "cell_2_uri",
-                        "languageId": "python",
-                        "text": "",
-                    },
-                ],
-            },
-        )
+        send_notebook_did_open(client, ["import sys", ""])
         wait_for_condition(lambda: mock_notify.call_count >= 2)
         assert len(server.workspace.documents) == 3
         for uri in ["cell_1_uri", "cell_2_uri", "notebook_uri"]:
@@ -531,48 +423,11 @@ def test_notebook__did_close(
     client_server_pair,
 ):
     client, server = client_server_pair
-    client._endpoint.request(
-        "initialize",
-        {
-            "processId": 1234,
-            "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
-        },
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    send_initialize_request(client)
 
     # Open notebook
     with patch.object(server._endpoint, "notify") as mock_notify:
-        client._endpoint.notify(
-            "notebookDocument/didOpen",
-            {
-                "notebookDocument": {
-                    "uri": "notebook_uri",
-                    "notebookType": "jupyter-notebook",
-                    "cells": [
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_1_uri",
-                        },
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_2_uri",
-                        },
-                    ],
-                },
-                "cellTextDocuments": [
-                    {
-                        "uri": "cell_1_uri",
-                        "languageId": "python",
-                        "text": "import sys",
-                    },
-                    {
-                        "uri": "cell_2_uri",
-                        "languageId": "python",
-                        "text": "",
-                    },
-                ],
-            },
-        )
+        send_notebook_did_open(client, ["import sys", ""])
         wait_for_condition(lambda: mock_notify.call_count >= 2)
         assert len(server.workspace.documents) == 3
         for uri in ["cell_1_uri", "cell_2_uri", "notebook_uri"]:
@@ -603,48 +458,11 @@ def test_notebook__did_close(
 @pytest.mark.skipif(IS_WIN, reason="Flaky on Windows")
 def test_notebook_definition(client_server_pair):
     client, server = client_server_pair
-    client._endpoint.request(
-        "initialize",
-        {
-            "processId": 1234,
-            "rootPath": os.path.dirname(__file__),
-            "initializationOptions": {},
-        },
-    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+    send_initialize_request(client)
 
     # Open notebook
     with patch.object(server._endpoint, "notify") as mock_notify:
-        client._endpoint.notify(
-            "notebookDocument/didOpen",
-            {
-                "notebookDocument": {
-                    "uri": "notebook_uri",
-                    "notebookType": "jupyter-notebook",
-                    "cells": [
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_1_uri",
-                        },
-                        {
-                            "kind": NotebookCellKind.Code,
-                            "document": "cell_2_uri",
-                        },
-                    ],
-                },
-                "cellTextDocuments": [
-                    {
-                        "uri": "cell_1_uri",
-                        "languageId": "python",
-                        "text": "y=2\nx=1",
-                    },
-                    {
-                        "uri": "cell_2_uri",
-                        "languageId": "python",
-                        "text": "x",
-                    },
-                ],
-            },
-        )
+        send_notebook_did_open(client, ["y=2\nx=1", "x"])
         # wait for expected diagnostics messages
         wait_for_condition(lambda: mock_notify.call_count >= 2)
         assert len(server.workspace.documents) == 3
diff --git a/test/test_utils.py b/test/test_utils.py
index bc2782d..fb4a8b8 100644
--- a/test/test_utils.py
+++ b/test/test_utils.py
@@ -6,15 +6,72 @@ import os
 import sys
 from threading import Thread
 import time
+from typing import List
 from unittest import mock
 
 from flaky import flaky
 from docstring_to_markdown import UnknownFormatError
 
 from pylsp import _utils
+from pylsp.lsp import NotebookCellKind
 from pylsp.python_lsp import PythonLSPServer, start_io_lang_server
 
 
+CALL_TIMEOUT_IN_SECONDS = 30
+
+
+def send_notebook_did_open(client, cells: List[str]):
+    """
+    Sends a notebookDocument/didOpen notification with the given python cells.
+
+    The notebook has the uri "notebook_uri" and the cells have the uris
+    "cell_1_uri", "cell_2_uri", etc.
+    """
+    client._endpoint.notify(
+        "notebookDocument/didOpen", notebook_with_python_cells(cells)
+    )
+
+
+def notebook_with_python_cells(cells: List[str]):
+    """
+    Create a notebook document with the given python cells.
+
+    The notebook has the uri "notebook_uri" and the cells have the uris
+    "cell_1_uri", "cell_2_uri", etc.
+    """
+    return {
+        "notebookDocument": {
+            "uri": "notebook_uri",
+            "notebookType": "jupyter-notebook",
+            "cells": [
+                {
+                    "kind": NotebookCellKind.Code,
+                    "document": f"cell_{i+1}_uri",
+                }
+                for i in range(len(cells))
+            ],
+        },
+        "cellTextDocuments": [
+            {
+                "uri": f"cell_{i+1}_uri",
+                "languageId": "python",
+                "text": cell,
+            }
+            for i, cell in enumerate(cells)
+        ],
+    }
+
+
+def send_initialize_request(client):
+    return client._endpoint.request(
+        "initialize",
+        {
+            "processId": 1234,
+            "rootPath": os.path.dirname(__file__),
+        },
+    ).result(timeout=CALL_TIMEOUT_IN_SECONDS)
+
+
 def start(obj):
     obj.start()