Fix gds integration test + some test refactoring

.github/workflows/gds-integration-tests.yml

@@ -39,5 +39,5 @@ jobs:
         env:
           AURA_API_CLIENT_ID: 4V1HYCYEeoU4dSxThKnBeLvE2U4hSphx
           AURA_API_CLIENT_SECRET: ${{ secrets.AURA_API_CLIENT_SECRET }}
-          AURA_API_TENANT_ID: 3f8df5e7-4800-4d4f-ad1d-2d044dfd587c
+          AURA_API_PROJECT_ID: 3f8df5e7-4800-4d4f-ad1d-2d044dfd587c
         run: uv run pytest tests/ --include-neo4j-and-gds

python-wrapper/pyproject.toml

@@ -116,8 +116,10 @@ filterwarnings = [
     "error",
     "ignore:Jupyter is migrating its paths to use standard platformdirs:DeprecationWarning",
     # snowflake vendors an older `requests` whose dependency check rejects the chardet
-    # version pulled in transitively by the notebook group. Harmless; ignore it.
-    "ignore:.*doesn't match a supported version:snowflake.connector.vendored.requests.exceptions.RequestsDependencyWarning"
+    # version (pulled in transitively via `encutils`). Harmless; ignore it.
+    # Matched by message only: naming the category here would make pytest import
+    # `snowflake.connector` to resolve it, and that import is exactly what fails.
+    "ignore:.*doesn't match a supported version"
 ]
 
 [tool.ruff]

python-wrapper/tests/conftest.py

@@ -1,11 +1,7 @@
-import os
-import random
-from typing import Any, Generator
-
 import pytest
 
 
-def pytest_addoption(parser: Any) -> None:
+def pytest_addoption(parser: pytest.Parser) -> None:
     parser.addoption(
         "--include-neo4j-and-gds",
         action="store_true",
@@ -18,7 +14,7 @@ def pytest_addoption(parser: Any) -> None:
     )
 
 
-def pytest_collection_modifyitems(config: Any, items: Any) -> None:
+def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
     if not config.getoption("--include-neo4j-and-gds"):
         skip = pytest.mark.skip(reason="skipping since requiring Neo4j instance with GDS running")
         for item in items:
@@ -29,96 +25,3 @@ def pytest_collection_modifyitems(config: Any, items: Any) -> None:
         for item in items:
             if "requires_snowflake" in item.keywords:
                 item.add_marker(skip)
-
-
-@pytest.fixture(scope="package")
-def aura_db_instance() -> Generator[Any, None, None]:
-    if os.environ.get("NEO4J_URI", ""):
-        print(f"Skipping Aura DB setup since NEO4J_URI is set to {os.environ['NEO4J_URI']}")
-        yield None
-        return
-
-    if os.environ.get("AURA_API_CLIENT_ID", None) is None:
-        yield None
-        return
-
-    from tests.gds_helper import aura_api, create_auradb_instance, wait_for_instance
-
-    api = aura_api()
-    instance_details = create_auradb_instance(api)
-
-    try:
-        dbms_connection_info = wait_for_instance(api, instance_details)
-
-        # setting as environment variables to run notebooks with this connection
-        os.environ["NEO4J_URI"] = dbms_connection_info.get_uri()
-        assert isinstance(dbms_connection_info.username, str)
-        os.environ["NEO4J_USERNAME"] = dbms_connection_info.username
-        assert isinstance(dbms_connection_info.password, str)
-        os.environ["NEO4J_PASSWORD"] = dbms_connection_info.password
-        old_instance = os.environ.get("AURA_INSTANCEID", "")
-        if dbms_connection_info.aura_instance_id:
-            os.environ["AURA_INSTANCEID"] = dbms_connection_info.aura_instance_id
-
-        yield dbms_connection_info
-
-        # Clear Neo4j_URI after test (rerun should create a new instance)
-        os.environ["AURA_INSTANCEID"] = old_instance
-        assert dbms_connection_info.aura_instance_id is not None
-    finally:
-        api.delete_instance(instance_details.id)
-
-
-@pytest.fixture(scope="package")
-def gds(aura_db_instance: Any) -> Generator[Any, None, None]:
-    from graphdatascience.session import SessionMemory
-
-    from tests.gds_helper import connect_to_local_gds_session, connect_to_plugin_gds, gds_sessions
-
-    if aura_db_instance:
-        sessions = gds_sessions()
-
-        gds = sessions.get_or_create(
-            f"neo4j-viz-ci-{os.environ.get('GITHUB_RUN_ID', random.randint(0, 10**6))}",
-            memory=SessionMemory.m_2GB,
-            db_connection=aura_db_instance,
-        )
-
-        yield gds
-        gds.delete()
-    else:
-        neo4j_uri = os.environ["NEO4J_URI"]
-        neo4j_auth = (os.environ.get("NEO4J_USERNAME", "neo4j"), os.environ.get("NEO4J_PASSWORD", "password"))
-
-        session_uri = os.environ.get("GDS_SESSION_URI")
-        if session_uri:
-            gds = connect_to_local_gds_session(session_uri, neo4j_uri, neo4j_auth)
-        else:
-            gds = connect_to_plugin_gds(neo4j_uri, neo4j_auth)  # type: ignore
-        yield gds
-        gds.close()
-
-
-@pytest.fixture(scope="package")
-def neo4j_driver(aura_db_instance: Any) -> Generator[Any, None, None]:
-    import neo4j
-
-    if aura_db_instance:
-        driver = neo4j.GraphDatabase.driver(
-            aura_db_instance.uri, auth=(aura_db_instance.username, aura_db_instance.password)
-        )
-    else:
-        NEO4J_URI = os.environ.get("NEO4J_URI", "neo4j://localhost:7687")
-        driver = neo4j.GraphDatabase.driver(NEO4J_URI)
-
-    try:
-        driver.verify_connectivity()
-        yield driver
-    finally:
-        driver.close()
-
-
-@pytest.fixture(scope="package")
-def neo4j_session(neo4j_driver: Any) -> Generator[Any, None, None]:
-    with neo4j_driver.session() as session:
-        yield session

python-wrapper/tests/neo4j_and_gds/conftest.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import os
+import random
+from typing import Generator
+
+import pytest
+
+pytest.importorskip("graphdatascience")
+
+import neo4j
+from graphdatascience import GraphDataScience
+from graphdatascience.session import AuraGraphDataScience, DbmsConnectionInfo, SessionMemory
+
+from tests.neo4j_and_gds.gds_helper import (
+    aura_api,
+    connect_to_local_gds_session,
+    connect_to_plugin_gds,
+    create_auradb_instance,
+    gds_sessions,
+    wait_for_instance,
+)
+
+
+@pytest.fixture(scope="package")
+def aura_db_instance() -> Generator[DbmsConnectionInfo | None, None, None]:
+    if os.environ.get("NEO4J_URI", ""):
+        print(f"Skipping Aura DB setup since NEO4J_URI is set to {os.environ['NEO4J_URI']}")
+        yield None
+        return
+
+    if os.environ.get("AURA_API_CLIENT_ID", None) is None:
+        yield None
+        return
+
+    api = aura_api()
+    instance_details = create_auradb_instance(api)
+
+    try:
+        dbms_connection_info = wait_for_instance(api, instance_details)
+
+        # setting as environment variables to run notebooks with this connection
+        os.environ["NEO4J_URI"] = dbms_connection_info.get_uri()
+        assert isinstance(dbms_connection_info.username, str)
+        os.environ["NEO4J_USERNAME"] = dbms_connection_info.username
+        assert isinstance(dbms_connection_info.password, str)
+        os.environ["NEO4J_PASSWORD"] = dbms_connection_info.password
+        old_instance = os.environ.get("AURA_INSTANCEID", "")
+        if dbms_connection_info.aura_instance_id:
+            os.environ["AURA_INSTANCEID"] = dbms_connection_info.aura_instance_id
+
+        yield dbms_connection_info
+
+        # Clear Neo4j_URI after test (rerun should create a new instance)
+        os.environ["AURA_INSTANCEID"] = old_instance
+        assert dbms_connection_info.aura_instance_id is not None
+    finally:
+        api.delete_instance(instance_details.id)
+
+
+@pytest.fixture(scope="package")
+def gds(aura_db_instance: DbmsConnectionInfo | None) -> Generator[GraphDataScience | AuraGraphDataScience, None, None]:
+    if aura_db_instance:
+        sessions = gds_sessions()
+
+        gds = sessions.get_or_create(
+            f"neo4j-viz-ci-{os.environ.get('GITHUB_RUN_ID', random.randint(0, 10**6))}",
+            memory=SessionMemory.m_2GB,
+            db_connection=aura_db_instance,
+        )
+
+        yield gds
+        gds.delete()
+    else:
+        neo4j_uri = os.environ["NEO4J_URI"]
+        neo4j_auth = (os.environ.get("NEO4J_USERNAME", "neo4j"), os.environ.get("NEO4J_PASSWORD", "password"))
+
+        session_uri = os.environ.get("GDS_SESSION_URI")
+        if session_uri:
+            gds = connect_to_local_gds_session(session_uri, neo4j_uri, neo4j_auth)
+        else:
+            gds = connect_to_plugin_gds(neo4j_uri, neo4j_auth)  # type: ignore
+        yield gds
+        gds.close()
+
+
+@pytest.fixture(scope="package")
+def neo4j_driver(aura_db_instance: DbmsConnectionInfo | None) -> Generator[neo4j.Driver, None, None]:
+    if aura_db_instance:
+        assert aura_db_instance.uri is not None
+        assert aura_db_instance.username is not None
+        assert aura_db_instance.password is not None
+        driver = neo4j.GraphDatabase.driver(
+            aura_db_instance.uri, auth=(aura_db_instance.username, aura_db_instance.password)
+        )
+    else:
+        NEO4J_URI = os.environ.get("NEO4J_URI", "neo4j://localhost:7687")
+        driver = neo4j.GraphDatabase.driver(NEO4J_URI)
+
+    try:
+        driver.verify_connectivity()
+        yield driver
+    finally:
+        driver.close()
+
+
+@pytest.fixture(scope="package")
+def neo4j_session(neo4j_driver: neo4j.Driver) -> Generator[neo4j.Session, None, None]:
+    with neo4j_driver.session() as session:
+        yield session

python-wrapper/tests/test_gds.py → python-wrapper/tests/neo4j_and_gds/test_gds.py

@@ -1,14 +1,19 @@
 import re
-from typing import Any, Generator
+from contextlib import AbstractContextManager
+from typing import Generator
 
 import pandas as pd
 import pytest
+from graphdatascience import GraphDataScience
+from graphdatascience.graph.v2 import GraphV2
+from graphdatascience.session import AuraGraphDataScience
 
 from neo4j_viz import Node
+from neo4j_viz.gds import from_gds
 
 
 @pytest.fixture(scope="class")
-def db_setup(gds: Any) -> Generator[None, None, None]:
+def db_setup(gds: GraphDataScience | AuraGraphDataScience) -> Generator[None, None, None]:
     gds.run_cypher(
         "CREATE "
         "  (a:_CI_A {name:'Alice', height:20, id:42, _id: 1337, caption: 'hello'})"
@@ -20,10 +25,7 @@ def db_setup(gds: Any) -> Generator[None, None, None]:
     gds.run_cypher("MATCH (n:_CI_A|_CI_B) DETACH DELETE n")
 
 
-def project_graph(gds: Any) -> Any:
-    from graphdatascience import GraphDataScience
-    from graphdatascience.session import AuraGraphDataScience
-
+def project_graph(gds: GraphDataScience | AuraGraphDataScience) -> AbstractContextManager[GraphV2]:
     if isinstance(gds, GraphDataScience):
         return gds.v2.graph.project("g2", "*", "*")
     elif isinstance(gds, AuraGraphDataScience):
@@ -33,9 +35,7 @@ def project_graph(gds: Any) -> Any:
 
 @pytest.mark.filterwarnings("ignore::DeprecationWarning")
 @pytest.mark.requires_neo4j_and_gds
-def test_from_gds_integration_all_db_properties(gds: Any, db_setup: None) -> None:
-    from neo4j_viz.gds import from_gds
-
+def test_from_gds_integration_all_db_properties(gds: GraphDataScience | AuraGraphDataScience, db_setup: None) -> None:
     with project_graph(gds) as G:
         VG = from_gds(gds, G, db_node_properties=["name"])
 
@@ -44,9 +44,7 @@ def test_from_gds_integration_all_db_properties(gds: Any, db_setup: None) -> Non
 
 
 @pytest.mark.requires_neo4j_and_gds
-def test_from_gds_integration_all_properties(gds: Any) -> None:
-    from neo4j_viz.gds import from_gds
-
+def test_from_gds_integration_all_properties(gds: GraphDataScience | AuraGraphDataScience) -> None:
     nodes = pd.DataFrame(
         {
             "nodeId": [0, 1, 2],
@@ -114,9 +112,7 @@ def test_from_gds_integration_all_properties(gds: Any) -> None:
 
 
 @pytest.mark.requires_neo4j_and_gds
-def test_from_gds_sample(gds: Any) -> None:
-    from neo4j_viz.gds import from_gds
-
+def test_from_gds_sample(gds: GraphDataScience | AuraGraphDataScience) -> None:
     with gds.v2.graph.generate("hello", node_count=11_000, average_degree=1) as G:
         with pytest.warns(
             UserWarning,
@@ -136,9 +132,7 @@ def test_from_gds_sample(gds: Any) -> None:
 
 
 @pytest.mark.requires_neo4j_and_gds
-def test_from_gds_hetero(gds: Any) -> None:
-    from neo4j_viz.gds import from_gds
-
+def test_from_gds_hetero(gds: GraphDataScience | AuraGraphDataScience) -> None:
     A_nodes = pd.DataFrame(
         {
             "nodeId": [0, 1],

python-wrapper/tests/neo4j_and_gds/test_notebooks.py

@@ -0,0 +1,17 @@
+import os
+
+import pytest
+from dotenv import load_dotenv
+from graphdatascience import GraphDataScience
+from graphdatascience.session import AuraGraphDataScience
+
+from tests.notebook_runner import run_notebooks
+
+
+@pytest.mark.requires_neo4j_and_gds
+def test_neo4j(gds: GraphDataScience | AuraGraphDataScience) -> None:
+    # The `gds` fixture provisions the Aura DB / GDS session and sets the NEO4J_* env vars
+    # that the notebooks read to connect.
+    load_dotenv(os.environ.get("ENV_FILE"))
+
+    run_notebooks(["neo4j-example.ipynb", "gds-example.ipynb"])