diff --git a/pylsp/python_lsp.py b/pylsp/python_lsp.py
index da0fc5fd..ce5f243d 100644
--- a/pylsp/python_lsp.py
+++ b/pylsp/python_lsp.py
@@ -634,8 +634,7 @@ def _cell_document__definition(self, cellDocument=None, position=None, **_kwargs
         with workspace.temp_document(total_source) as random_uri:
             definitions = self.definitions(random_uri, position)
             for definition in definitions:
-                # TODO: a better test for if a definition is the random_uri
-                if random_uri in definition['uri']:
+                if random_uri == definition['uri']:
                     # Find the cell the start is in
                     for cell in cell_list:
                         # TODO: perhaps it is more correct to check definition['range']['end']['line'] <= cell['line_end'], but 
diff --git a/pylsp/workspace.py b/pylsp/workspace.py
index a53cb345..1f084b4f 100644
--- a/pylsp/workspace.py
+++ b/pylsp/workspace.py
@@ -120,15 +120,15 @@ def put_notebook_document(self, doc_uri, notebook_type, cells, version=None, met
         self._docs[doc_uri] = self._create_notebook_document(doc_uri, notebook_type, cells, version, metadata)
 
     @contextmanager
-    def temp_document(self, source):
-        random_uri = str(uuid.uuid4())
+    def temp_document(self, source, path=None):
+        if path is None:
+            path = self.root_path
+        uri = uris.from_fs_path(os.path.join(path, str(uuid.uuid4())))
         try:
-            self.put_document(random_uri, source)
-            yield random_uri
+            self.put_document(uri, source)
+            yield uri
         finally:
-            self.rm_document(random_uri)
-
-
+            self.rm_document(uri)
 
     def add_notebook_cells(self, doc_uri, cells, start):
         self._docs[doc_uri].add_cells(cells, start)