name: Reusable Graph DB Tests

permissions:
  contents: read
  packages: read

on:
  workflow_call:
    inputs:
      databases:
        required: false
        type: string
        default: "all"
        description: "Which vector databases to test (comma-separated list or 'all')"
      python-version:
        required: false
        type: string
        default: "3.11.x"
        description: "Python version to use for testing"
      ci-image:
        required: false
        type: string
        default: ''

env:
  COGNEE_SKIP_CONNECTION_TEST: 'true'

jobs:
  run-kuzu-tests:
    name: Kuzu Tests
    runs-on: ubuntu-22.04
    container: ${{ inputs.ci-image != '' && fromJSON(format('{{"image":"{0}","credentials":{{"username":"{1}","password":"{2}"}}}}', inputs.ci-image, github.actor, github.token)) || null }}
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'kuzu') }}
    steps:
      - name: Check out
        uses: actions/checkout@v6
        with:
          fetch-depth: 0

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: ${{ inputs.python-version }}

      - name: Dependencies already installed
        run: echo "Dependencies already installed in setup"

      - name: Run Kuzu Tests
        env:
          ENV: 'dev'
          LLM_MODEL: ${{ secrets.LLM_MODEL }}
          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
          LLM_ARGS: ${{ secrets.LLM_ARGS }}
          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
          EMBEDDING_DIMENSIONS: 300
          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
        run: uv run python ./cognee/tests/test_kuzu.py

      - name: Run Weighted Edges Tests with Kuzu
        env:
          ENV: 'dev'
          GRAPH_DATABASE_PROVIDER: "kuzu"
          LLM_MODEL: ${{ secrets.LLM_MODEL }}
          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
          LLM_ARGS: ${{ secrets.LLM_ARGS }}
          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
          EMBEDDING_DIMENSIONS: 300
          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
        run: uv run pytest cognee/tests/unit/interfaces/graph/test_weighted_edges.py -v

  run-postgres-tests:
    name: Postgres Graph Tests
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'postgres') }}
    services:
      postgres:
        image: ghcr.io/topoteretes/pgvector:pg17
        credentials:
          username: ${{ github.actor }}
          password: ${{ secrets.GITHUB_TOKEN }}
        env:
          POSTGRES_USER: cognee
          POSTGRES_PASSWORD: cognee
          POSTGRES_DB: cognee_db
        ports:
          - 5432:5432
        options: >-
          --health-cmd pg_isready
          --health-interval 10s
          --health-timeout 5s
          --health-retries 5
    steps:
      - name: Check out
        uses: actions/checkout@v6
        with:
          fetch-depth: 0

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: ${{ inputs.python-version }}
          extra-dependencies: postgres

      - name: Run Postgres Adapter Unit Tests
        env:
          DB_PROVIDER: postgres
          DB_HOST: localhost
          DB_PORT: 5432
          DB_USERNAME: cognee
          DB_PASSWORD: cognee
          DB_NAME: cognee_db
        run: uv run pytest cognee/tests/e2e/postgres/test_postgres_adapter.py -v

      - name: Run Postgres Graph E2E Test
        env:
          ENV: 'dev'
          GRAPH_DATABASE_PROVIDER: "postgres"
          GRAPH_DATABASE_URL: "postgresql+asyncpg://cognee:cognee@localhost:5432/cognee_db"
          DB_PROVIDER: postgres
          DB_HOST: localhost
          DB_PORT: 5432
          DB_USERNAME: cognee
          DB_PASSWORD: cognee
          DB_NAME: cognee_db
          ENABLE_BACKEND_ACCESS_CONTROL: 'false'
          LLM_MODEL: ${{ secrets.LLM_MODEL }}
          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
          LLM_ARGS: ${{ secrets.LLM_ARGS }}
          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
          EMBEDDING_DIMENSIONS: 300
          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
        run: uv run python ./cognee/tests/e2e/postgres/test_graphdb_postgres.py

#      - name: Run Postgres Hybrid Adapter Tests
#        env:
#          ENV: 'dev'
#          DB_PROVIDER: postgres
#          DB_HOST: localhost
#          DB_PORT: 5432
#          DB_USERNAME: cognee
#          DB_PASSWORD: cognee
#          DB_NAME: cognee_db
#          EMBEDDING_DIMENSIONS: 300
#          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
#          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
#        run: uv run pytest cognee/tests/e2e/postgres/test_postgres_hybrid_adapter.py -v
#
#      - name: Run Postgres Hybrid E2E Test
#        env:
#          ENV: 'dev'
#          USE_UNIFIED_PROVIDER: "pghybrid"
#          DB_PROVIDER: postgres
#          DB_HOST: localhost
#          DB_PORT: 5432
#          DB_USERNAME: cognee
#          DB_PASSWORD: cognee
#          DB_NAME: cognee_db
#          ENABLE_BACKEND_ACCESS_CONTROL: 'false'
#          LLM_MODEL: ${{ secrets.LLM_MODEL }}
#          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
#          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
#          LLM_ARGS: ${{ secrets.LLM_ARGS }}
#          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
#          EMBEDDING_DIMENSIONS: 300
#          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
#          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
#        run: uv run python ./cognee/tests/e2e/postgres/test_pghybrid.py

  run-neo4j-tests:
    name: Neo4j Tests
    runs-on: ubuntu-22.04
    if: ${{ inputs.databases == 'all' || contains(inputs.databases, 'neo4j') }}
    steps:
      - name: Check out
        uses: actions/checkout@master

      - name: Cognee Setup
        uses: ./.github/actions/cognee_setup
        with:
          python-version: ${{ inputs.python-version }}

      - name: Setup Neo4j with GDS
        uses: ./.github/actions/setup_neo4j
        id: neo4j

      - name: Run default Neo4j
        env:
          ENV: 'dev'
          ENABLE_BACKEND_ACCESS_CONTROL: 'false'
          LLM_MODEL: ${{ secrets.LLM_MODEL }}
          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
          LLM_ARGS: ${{ secrets.LLM_ARGS }}
          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
          EMBEDDING_DIMENSIONS: 300
          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
          GRAPH_DATABASE_PROVIDER: "neo4j"
          GRAPH_DATABASE_URL: ${{ steps.neo4j.outputs.neo4j-url }}
          GRAPH_DATABASE_USERNAME: ${{ steps.neo4j.outputs.neo4j-username }}
          GRAPH_DATABASE_PASSWORD: ${{ steps.neo4j.outputs.neo4j-password }}
        run: uv run python ./cognee/tests/test_neo4j.py

      - name: Run Weighted Edges Tests with Neo4j
        env:
          ENV: 'dev'
          GRAPH_DATABASE_PROVIDER: "neo4j"
          GRAPH_DATABASE_URL: ${{ secrets.NEO4J_API_URL }}
          GRAPH_DATABASE_PASSWORD: ${{ secrets.NEO4J_API_KEY }}
          GRAPH_DATABASE_USERNAME: "neo4j"
          LLM_MODEL: ${{ secrets.LLM_MODEL }}
          LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
          LLM_API_KEY: ${{ secrets.LLM_API_KEY }}
          LLM_ARGS: ${{ secrets.LLM_ARGS }}
          LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
          EMBEDDING_DIMENSIONS: 300
          EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
          EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
        run: uv run pytest cognee/tests/unit/interfaces/graph/test_weighted_edges.py -v