test: Restore knowldege-store integration tests (#529)

* test: Restore knowldege-store integration tests

These currently just demonstrate that the knowledge store can be
created. That is enough to verify the schema and queries are correct.
I'll add more tests over time.

These are smaller than the full langchain tests to run, so are easier
to use for iterating on the details.

Author: bjchambers

.github/workflows/ci-unit-tests.yml

@@ -147,6 +147,9 @@ jobs:
           if [[ "true" == "${{ needs.preconditions.outputs.libs_colbert }}" ]]; then
             run_itests libs/colbert
           fi
+          if [[ "true" == "${{ needs.preconditions.outputs.libs_knowledge_store }}" ]]; then
+            run_itests libs/knowledge-store
+          fi
           if [[ "true" == "${{ needs.preconditions.outputs.libs_langchain }}" ]]; then
             run_itests libs/langchain
           fi

libs/knowledge-store/pyproject.toml

@@ -10,7 +10,7 @@ documentation = "https://docs.datastax.com/en/ragstack"
 packages = [{ include = "ragstack_knowledge_store" }]
 
 [tool.poetry.dependencies]
-python = ">=3.8.1,<4.0"
+python = ">=3.9,<3.13"
 cassio = "^0.1.7"
 
 [tool.poetry.group.dev.dependencies]
@@ -29,6 +29,10 @@ numpy = [
     {version = ">1.26.0,<2", python = ">=3.12"}
 ]
 
+[tool.poetry.group.test.dependencies]
+ragstack-ai-langchain = { path = "../langchain", develop = true }
+ragstack-ai-tests-utils = { path = "../tests-utils", develop = true }
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"

libs/knowledge-store/tests/integration_tests/test_graph_store.py

@@ -0,0 +1,54 @@
+import secrets
+from typing import Iterator
+
+import pytest
+from dotenv import load_dotenv
+from langchain_openai import OpenAIEmbeddings
+from ragstack_knowledge_store.graph_store import GraphStore
+from ragstack_tests_utils import LocalCassandraTestStore
+
+load_dotenv()
+
+KEYSPACE = "default_keyspace"
+
+
+@pytest.fixture(scope="session")
+def cassandra() -> Iterator[LocalCassandraTestStore]:
+    store = LocalCassandraTestStore()
+    yield store
+
+    if store.docker_container:
+        store.docker_container.stop()
+
+
+@pytest.fixture
+def graph_store_factory(cassandra: LocalCassandraTestStore):
+    session = cassandra.create_cassandra_session()
+    session.set_keyspace(KEYSPACE)
+
+    embedding = OpenAIEmbeddings()
+
+    def _make_graph_store():
+        name = secrets.token_hex(8)
+
+        node_table = f"nodes_{name}"
+        targets_table = f"targets_{name}"
+        return GraphStore(
+            embedding,
+            session=session,
+            keyspace=KEYSPACE,
+            node_table=node_table,
+            targets_table=targets_table,
+        )
+
+    yield _make_graph_store
+
+    session.shutdown()
+
+
+def test_graph_store_creation(graph_store_factory):
+    """Test that a graph store can be created.
+
+    This verifies the schema can be applied and the queries prepared.
+    """
+    graph_store_factory()

libs/knowledge-store/tox.ini

@@ -11,4 +11,15 @@ pass_env =
 commands =
     poetry install
     poetry build
-    poetry run pytest --disable-warnings {toxinidir}/tests
+    poetry run pytest --disable-warnings {toxinidir}/tests/unit_tests
+
+[testenv:integration-tests]
+description = run integration tests
+deps =
+    poetry
+pass_env =
+    OPENAI_API_KEY
+commands =
+    poetry install
+    poetry build
+    poetry run pytest --disable-warnings {toxinidir}/tests/integration_tests

libs/langchain/tests/integration_tests/conftest.py

@@ -26,6 +26,7 @@ def get_astradb_test_store() -> AstraDBTestStore:
 @pytest.fixture(scope="session", autouse=True)
 def before_after_tests():
     yield
+
     if (
         status["local_cassandra_test_store"]
         and status["local_cassandra_test_store"].docker_container

libs/tests-utils/ragstack_tests_utils/cassandra_container.py

@@ -12,6 +12,14 @@ def __init__(
         super().__init__(image=image, **kwargs)
         self.port = port
 
+        self.with_env(
+            "JVM_OPTS",
+            "-Dcassandra.skip_wait_for_gossip_to_settle=0 -Dcassandra.initial_token=0",
+        )
+        self.with_env("HEAP_NEWSIZE", "128M")
+        self.with_env("MAX_HEAP_SIZE", "1024M")
+        self.with_env("CASSANDRA_ENDPOINT_SNITCH", "GossipingPropertyFileSnitch")
+        self.with_env("CASSANDRA_DC", "datacenter1")
         self.with_exposed_ports(self.port)
 
     def _configure(self):