feature: add the test to workflows

This commit is contained in:
Hande 2025-11-04 16:45:46 +01:00
parent 75d705463f
commit b7047b0267
3 changed files with 24 additions and 66 deletions

View file

@ -259,3 +259,27 @@ jobs:
EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
run: uv run python ./cognee/tests/test_add_docling_document.py
test-custom-model:
name: Run Custom Model Test
runs-on: ubuntu-22.04
steps:
- name: Check out repository
uses: actions/checkout@v4
- name: Cognee Setup
uses: ./.github/actions/cognee_setup
with:
python-version: '3.11.x'
- name: Run Custom Model Test
env:
LLM_MODEL: ${{ secrets.LLM_MODEL }}
LLM_ENDPOINT: ${{ secrets.LLM_ENDPOINT }}
LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
LLM_API_VERSION: ${{ secrets.LLM_API_VERSION }}
EMBEDDING_MODEL: ${{ secrets.EMBEDDING_MODEL }}
EMBEDDING_ENDPOINT: ${{ secrets.EMBEDDING_ENDPOINT }}
EMBEDDING_API_KEY: ${{ secrets.EMBEDDING_API_KEY }}
EMBEDDING_API_VERSION: ${{ secrets.EMBEDDING_API_VERSION }}
run: uv run python ./cognee/tests/test_custom_model.py

0
cognee/tests/test_custom_model.py Executable file → Normal file
View file

View file

@ -1,66 +0,0 @@
import unittest
import subprocess
import os
import sys
class TestPipelines(unittest.TestCase):
"""Tests that all pipelines run successfully."""
def setUp(self):
# Ensure we're in the correct directory
self.project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
self.pipelines_dir = os.path.join(self.project_root, "src", "pipelines")
# Required environment variables
self.required_env_vars = ["LLM_API_KEY", "EMBEDDING_API_KEY"]
# Check if required environment variables are set
missing_vars = [var for var in self.required_env_vars if not os.environ.get(var)]
if missing_vars:
self.skipTest(f"Missing required environment variables: {', '.join(missing_vars)}")
def _run_pipeline(self, script_name):
"""Helper method to run a pipeline script and return the result."""
script_path = os.path.join(self.pipelines_dir, script_name)
# Use the Python executable from the virtual environment
python_exe = os.path.join(self.project_root, ".venv", "bin", "python")
if not os.path.exists(python_exe):
python_exe = sys.executable
try:
result = subprocess.run(
[python_exe, script_path],
check=True,
capture_output=True,
text=True,
timeout=300, # 5 minute timeout
)
return result
except subprocess.CalledProcessError as e:
self.fail(
f"Pipeline {script_name} failed with code {e.returncode}. "
f"Stdout: {e.stdout}, Stderr: {e.stderr}"
)
except subprocess.TimeoutExpired:
self.fail(f"Pipeline {script_name} timed out after 300 seconds")
def test_default_pipeline(self):
"""Test that the default pipeline runs successfully."""
result = self._run_pipeline("default.py")
self.assertEqual(result.returncode, 0)
def test_low_level_pipeline(self):
"""Test that the low-level pipeline runs successfully."""
result = self._run_pipeline("low_level.py")
self.assertEqual(result.returncode, 0)
def test_custom_model_pipeline(self):
"""Test that the custom model pipeline runs successfully."""
result = self._run_pipeline("custom-model.py")
self.assertEqual(result.returncode, 0)
if __name__ == "__main__":
unittest.main()