Updates tests to work with new Neo4j calendar versioning (#246)

Author: alexthomas93

tests/unit/experimental/components/test_kg_writer.py

@@ -14,6 +14,7 @@
 #  limitations under the License.
 from __future__ import annotations
 
+from typing import Tuple
 from unittest import mock
 from unittest.mock import MagicMock, Mock
 
@@ -331,3 +332,45 @@ async def test_run_is_version_5_23_or_above(_: Mock) -> None:
         parameters_=parameters_,
         database_=None,
     )
+
+
+@pytest.mark.parametrize(
+    "description, version, version_tuple, is_5_23_or_above",
+    [
+        ("SemVer, < 5.23", "5.22.0", (5, 22, 0), False),
+        ("SemVer, > 5.23", "5.24.0", (5, 24, 0), True),
+        ("SemVer, < 5.23, Aura", "5.22-aura", (5, 22, 0), False),
+        ("SemVer, > 5.23, Aura", "5.24-aura", (5, 24, 0), True),
+        ("CalVer", "2025.01.0", (2025, 1, 0), True),
+        ("CalVer, Aura", "2025.01-aura", (2025, 1, 0), True),
+    ],
+)
+@mock.patch(
+    "neo4j_graphrag.experimental.components.kg_writer.Neo4jWriter._db_setup",
+    return_value=None,
+)
+def test_get_version(
+    _: Mock,
+    driver: MagicMock,
+    description: str,
+    version: str,
+    version_tuple: Tuple[int],
+    is_5_23_or_above: bool,
+) -> None:
+    execute_query_mock = MagicMock(
+        return_value=(
+            [
+                {"versions": [version]},
+            ],
+            None,
+            None,
+        )
+    )
+    driver.execute_query = execute_query_mock
+    neo4j_writer = Neo4jWriter(driver=driver)
+    assert (
+        version_tuple == neo4j_writer._get_version()
+    ), f"Failed version_tuple test case: {description}"
+    assert (
+        neo4j_writer.is_version_5_23_or_above is is_5_23_or_above
+    ), f"Failed is_version_5_23_or_above test case: {description}"

tests/unit/retrievers/test_base.py

@@ -32,6 +32,8 @@
         (["5.19.0"], ((5, 19, 0), False)),
         (["4.3.5"], ((4, 3, 5), False)),
         (["5.23.0-6698"], ((5, 23, 0), False)),
+        (["2025.01.0"], ((2025, 1, 0), False)),
+        (["2025.01-aura"], ((2025, 1, 0), True)),
     ],
 )
 def test_retriever_get_version(
@@ -58,6 +60,8 @@ def get_search_results(self, *args: Any, **kwargs: Any) -> RawSearchResult:
         (((5, 19, 0), False), None),
         (((4, 3, 5), False), Neo4jVersionError),
         (((5, 23, 0), False), None),
+        (((2025, 1, 0), False), None),
+        (((2025, 1, 0), True), None),
     ],
 )
 @patch("neo4j_graphrag.retrievers.base.Retriever._get_version")