참고소스 수정본

참고/neo4j-graphrag-python-main/tests/unit/utils/test_file_handler.py

@@ -0,0 +1,293 @@
+#  Copyright (c) "Neo4j"
+#  Neo4j Sweden AB [https://neo4j.com]
+#  #
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  #
+#      https://www.apache.org/licenses/LICENSE-2.0
+#  #
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+from pathlib import Path
+from unittest.mock import patch, Mock, mock_open
+
+import pytest
+
+from neo4j_graphrag.utils.file_handler import FileHandler, FileFormat
+
+
+def test_file_handler_guess_format() -> None:
+    handler = FileHandler()
+
+    assert handler._guess_file_format(Path("file.json")) == FileFormat.JSON
+    assert handler._guess_file_format(Path("file.JSON")) == FileFormat.JSON
+    assert handler._guess_file_format(Path("file.yaml")) == FileFormat.YAML
+    assert handler._guess_file_format(Path("file.YAML")) == FileFormat.YAML
+    assert handler._guess_file_format(Path("file.yml")) == FileFormat.YAML
+    assert handler._guess_file_format(Path("file.YML")) == FileFormat.YAML
+    assert handler._guess_file_format(Path("file.txt")) is None
+
+
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._read_json")
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._check_file_exists")
+def test_file_handler_read_json_from_read_method_happy_path(
+    mock_file_exists: Mock,
+    mock_read_json: Mock,
+) -> None:
+    handler = FileHandler()
+
+    mock_file_exists.return_value = Path("file.json")
+    mock_read_json.return_value = {}
+    data = handler.read("file.json")
+    mock_read_json.assert_called_with(Path("file.json"))
+    assert data == {}
+
+    mock_file_exists.return_value = Path("file.JSON")
+    mock_read_json.return_value = {}
+    data = handler.read("file.JSON")
+    mock_read_json.assert_called_with(Path("file.JSON"))
+    assert data == {}
+
+
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._read_yaml")
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._check_file_exists")
+def test_file_handler_read_yaml_from_read_method_happy_path(
+    mock_file_exists: Mock,
+    mock_read_yaml: Mock,
+) -> None:
+    mock_file_exists.return_value = Path("file.yaml")
+    handler = FileHandler()
+    mock_read_yaml.return_value = {}
+    data = handler.read("file.yaml")
+    mock_read_yaml.assert_called_with(Path("file.yaml"))
+    assert data == {}
+
+    mock_file_exists.return_value = Path("file.yml")
+    mock_read_yaml.return_value = {}
+    data = handler.read("file.yml")
+    mock_read_yaml.assert_called_with(Path("file.yml"))
+    assert data == {}
+
+    mock_file_exists.return_value = Path("file.YAML")
+    mock_read_yaml.return_value = {}
+    data = handler.read("file.YAML")
+    mock_read_yaml.assert_called_with(Path("file.YAML"))
+    assert data == {}
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+def test_file_handler_read_json_method_happy_path(
+    mock_fs: Mock,
+) -> None:
+    mock_fs_open = mock_open(read_data='{"data": 1}')
+    mock_fs.return_value.open = mock_fs_open
+
+    handler = FileHandler()
+    data = handler._read_json(Path("file.json"))
+    mock_fs_open.assert_called_once_with("file.json", "r")
+    assert data == {"data": 1}
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+def test_file_handler_read_yaml_method_happy_path(
+    mock_fs: Mock,
+) -> None:
+    mock_fs_open = mock_open(
+        read_data="""
+    data: 1
+    """
+    )
+    mock_fs.return_value.open = mock_fs_open
+
+    handler = FileHandler()
+    data = handler._read_yaml(Path("file.yaml"))
+    mock_fs_open.assert_called_once_with("file.yaml", "r")
+    assert data == {"data": 1}
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+def test_file_handler_read_json_file_does_not_exist(
+    mock_fs: Mock,
+) -> None:
+    mock_fs.return_value.open.side_effect = FileNotFoundError
+    handler = FileHandler()
+    with pytest.raises(FileNotFoundError):
+        handler._read_json(Path("file.json"))
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+def test_file_handler_read_json_invalid_json(
+    mock_fs: Mock,
+) -> None:
+    mock_fs_open = mock_open(read_data="{")
+    mock_fs.return_value.open = mock_fs_open
+
+    handler = FileHandler()
+    with pytest.raises(ValueError, match="Invalid JSON"):
+        handler._read_json(Path("file.json"))
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+def test_file_handler_read_yaml_file_does_not_exist(
+    mock_fs: Mock,
+) -> None:
+    mock_fs.return_value.open.side_effect = FileNotFoundError
+    handler = FileHandler()
+    with pytest.raises(FileNotFoundError):
+        handler._read_yaml(Path("file.yaml"))
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+def test_file_handler_read_yaml_invalid_yaml(
+    mock_fs: Mock,
+) -> None:
+    mock_fs_open = mock_open(
+        read_data="""
+    data: [
+    """
+    )
+    mock_fs.return_value.open = mock_fs_open
+
+    handler = FileHandler()
+    with pytest.raises(ValueError, match="Invalid YAML"):
+        handler._read_yaml(Path("file.yaml"))
+
+
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._check_file_exists")
+def test_file_handler_file_can_be_written_file(mock_file_exists: Mock) -> None:
+    # file does not exist
+    mock_file_exists.side_effect = FileNotFoundError()
+    handler = FileHandler()
+    # nothing happens = all good
+    handler._check_file_can_be_written(Path("file.json"))
+    # file exist, overwrite is False
+    mock_file_exists.side_effect = None
+    handler = FileHandler()
+    with pytest.raises(ValueError):
+        handler._check_file_can_be_written(Path("file.json"), overwrite=False)
+    # file exists, overwrite is True
+    mock_file_exists.side_effect = None
+    handler = FileHandler()
+    # nothing happens = all good
+    handler._check_file_can_be_written(Path("file.json"), overwrite=True)
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+@patch("neo4j_graphrag.utils.file_handler.json")
+def test_file_handler_write_json_happy_path(
+    mock_json_module: Mock,
+    mock_fs: Mock,
+) -> None:
+    mock_fs_open = mock_open()
+    mock_fs.return_value.open = mock_fs_open
+    file_handler = FileHandler()
+    file_handler._write_json({"some": "data"}, Path("file.json"))
+    mock_fs_open.assert_called_once_with("file.json", "w")
+    mock_json_module.dump.assert_called_with(
+        {"some": "data"}, mock_fs_open.return_value, indent=2
+    )
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+@patch("neo4j_graphrag.utils.file_handler.json")
+def test_file_handler_write_json_extra_kwargs_happy_path(
+    mock_json_module: Mock,
+    mock_fs: Mock,
+) -> None:
+    mock_fs_open = mock_open()
+    mock_fs.return_value.open = mock_fs_open
+    file_handler = FileHandler()
+    file_handler._write_json({"some": "data"}, Path("file.json"), indent=4, default=str)
+    mock_fs_open.assert_called_once_with("file.json", "w")
+    mock_json_module.dump.assert_called_with(
+        {"some": "data"}, mock_fs_open.return_value, indent=4, default=str
+    )
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+@patch("neo4j_graphrag.utils.file_handler.yaml")
+def test_file_handler_write_yaml_happy_path(
+    mock_yaml_module: Mock,
+    mock_fs: Mock,
+) -> None:
+    mock_fs_open = mock_open()
+    mock_fs.return_value.open = mock_fs_open
+    file_handler = FileHandler()
+    file_handler._write_yaml({"some": "data"}, Path("file.yaml"))
+    mock_fs_open.assert_called_once_with("file.yaml", "w")
+    mock_yaml_module.safe_dump.assert_called_with(
+        {"some": "data"},
+        mock_fs_open.return_value,
+        default_flow_style=False,
+        sort_keys=True,
+    )
+
+
+@patch("neo4j_graphrag.utils.file_handler.LocalFileSystem")
+@patch("neo4j_graphrag.utils.file_handler.yaml")
+def test_file_handler_write_yaml_extra_kwargs_happy_path(
+    mock_yaml_module: Mock,
+    mock_fs: Mock,
+) -> None:
+    mock_fs_open = mock_open()
+    mock_fs.return_value.open = mock_fs_open
+    file_handler = FileHandler()
+    file_handler._write_yaml(
+        {"some": "data"}, Path("file.json"), default_flow_style="toto", other_keyword=42
+    )
+    mock_fs_open.assert_called_once_with("file.json", "w")
+    mock_yaml_module.safe_dump.assert_called_with(
+        {"some": "data"},
+        mock_fs_open.return_value,
+        default_flow_style="toto",
+        sort_keys=True,
+        other_keyword=42,
+    )
+
+
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._write_json")
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._check_file_can_be_written")
+def test_file_handler_write_json_from_write_method_happy_path(
+    mock_file_can_be_written: Mock,
+    mock_write_json: Mock,
+) -> None:
+    handler = FileHandler()
+    handler.write("data", "file.json")
+    mock_write_json.assert_called_with("data", Path("file.json"))
+
+
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._write_yaml")
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._check_file_can_be_written")
+def test_file_handler_write_yaml_from_write_method_happy_path(
+    mock_file_can_be_written: Mock,
+    mock_write_yaml: Mock,
+) -> None:
+    handler = FileHandler()
+    handler.write("data", "file.yaml")
+    mock_write_yaml.assert_called_with("data", Path("file.yaml"))
+
+
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._write_yaml")
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._check_file_can_be_written")
+def test_file_handler_write_yaml_from_write_method_overwrite_format_happy_path(
+    mock_file_can_be_written: Mock,
+    mock_write_yaml: Mock,
+) -> None:
+    handler = FileHandler()
+    handler.write("data", "file.txt", format=FileFormat.YAML)
+    mock_write_yaml.assert_called_with("data", Path("file.txt"))
+
+
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._write_json")
+@patch("neo4j_graphrag.utils.file_handler.FileHandler._check_file_can_be_written")
+def test_file_handler_write_json_from_write_method_overwrite_format_happy_path(
+    mock_file_can_be_written: Mock,
+    mock_write_json: Mock,
+) -> None:
+    handler = FileHandler()
+    handler.write("data", "file.txt", format=FileFormat.JSON)
+    mock_write_json.assert_called_with("data", Path("file.txt"))