Moved output index logic to a separate method, added tests

Author: 3coins

jupyter_server_documents/outputs/manager.py

@@ -34,6 +34,30 @@ def _build_path(self, file_id, cell_id=None, output_index=None):
         if output_index is not None:
             path = path / f"{output_index}.output"
         return path
+    
+    def _determine_output_index(self, cell_id, display_id=None):
+        """
+        Determines output index for a given cell and optional display ID.
+        
+        Args:
+            cell_id (str): The cell identifier
+            display_id (str, optional): A display identifier. Defaults to None.
+        
+        Returns:
+            int: The allocated output index
+        """
+        last_index = self._last_output_index.get(cell_id, -1)
+        if display_id:
+            index = self._output_index_by_display_id.get(display_id)
+            if index is None:
+                index = last_index + 1
+                self._last_output_index[cell_id] = index
+                self._output_index_by_display_id[display_id] = index
+        else:
+            index = last_index + 1
+            self._last_output_index[cell_id] = index
+        
+        return index
 
     def get_output_index(self, display_id: str):
         """Returns output index for a cell by display_id"""
@@ -121,17 +145,7 @@ def write(self, file_id, cell_id, output, display_id=None):
 
     def write_output(self, file_id, cell_id, output, display_id=None):
         self._ensure_path(file_id, cell_id)
-        last_index = self._last_output_index.get(cell_id, -1)
-        if display_id:
-            index = self._output_index_by_display_id.get(display_id)
-            if index is None:
-                index = last_index + 1
-                self._last_output_index[cell_id] = index
-                self._output_index_by_display_id[display_id] = index
-        else:
-            index = last_index + 1
-            self._last_output_index[cell_id] = index
-        
+        index = self._determine_output_index(cell_id, display_id)
         path = self._build_path(file_id, cell_id, index)
         data = json.dumps(output, ensure_ascii=False)
         with open(path, "w", encoding="utf-8") as f:

jupyter_server_documents/tests/test_outputs_manager.py

@@ -91,3 +91,41 @@ def file_not_found():
             op.get_output('a','b',0)
         with pytest.raises(FileNotFoundError):
             op.get_stream('a','b')       
+
+
+@pytest.mark.parametrize(
+    "_last_output_index, _output_index_by_display_id, cell_id, display_id, expected", [
+    # Empty dictionaries, no display_id
+    ({}, {}, "cell1", None, 0),
+    
+    # Empty dictionaries, with display_id
+    ({}, {}, "cell1", "display1", 0),
+    
+    # Existing last_output_index, no display_id
+    ({"cell1": 5}, {}, "cell1", None, 6),
+    ({"cell2": 3}, {}, "cell1", None, 0),
+    
+    # Existing output_index_by_display_id
+    ({}, {"display1": 2}, "cell1", "display1", 2),
+    
+    # Existing last_output_index and display_id
+    ({"cell1": 5}, {"display1": 3}, "cell1", "display1", 3),
+    ({"cell1": 5}, {"display1": 3}, "cell1", "display2", 6),
+    
+    # Multiple cells with different indices
+    ({"cell1": 2, "cell2": 7}, {}, "cell1", None, 3),
+    ({"cell1": 2, "cell2": 7}, {}, "cell3", None, 0),
+])
+def test__determine_output_index(
+    _last_output_index, 
+    _output_index_by_display_id, 
+    cell_id, 
+    display_id,
+    expected
+):
+    op = OutputsManager()
+    op._last_output_index = _last_output_index;
+    op._output_index_by_display_id = _output_index_by_display_id
+    
+    assert expected == op._determine_output_index(cell_id, display_id)
+