Feat add testcontainers test (#23269)
This commit is contained in:
3
.github/workflows/api-tests.yml
vendored
3
.github/workflows/api-tests.yml
vendored
@@ -99,3 +99,6 @@ jobs:
|
||||
|
||||
- name: Run Tool
|
||||
run: uv run --project api bash dev/pytest/pytest_tools.sh
|
||||
|
||||
- name: Run TestContainers
|
||||
run: uv run --project api bash dev/pytest/pytest_testcontainers.sh
|
||||
|
@@ -114,6 +114,7 @@ dev = [
|
||||
"pytest-cov~=4.1.0",
|
||||
"pytest-env~=1.1.3",
|
||||
"pytest-mock~=3.14.0",
|
||||
"testcontainers~=4.10.0",
|
||||
"types-aiofiles~=24.1.0",
|
||||
"types-beautifulsoup4~=4.12.0",
|
||||
"types-cachetools~=5.5.0",
|
||||
|
328
api/tests/test_containers_integration_tests/conftest.py
Normal file
328
api/tests/test_containers_integration_tests/conftest.py
Normal file
@@ -0,0 +1,328 @@
|
||||
"""
|
||||
TestContainers-based integration test configuration for Dify API.
|
||||
|
||||
This module provides containerized test infrastructure using TestContainers library
|
||||
to spin up real database and service instances for integration testing. This approach
|
||||
ensures tests run against actual service implementations rather than mocks, providing
|
||||
more reliable and realistic test scenarios.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from flask.testing import FlaskClient
|
||||
from sqlalchemy.orm import Session
|
||||
from testcontainers.core.container import DockerContainer
|
||||
from testcontainers.core.waiting_utils import wait_for_logs
|
||||
from testcontainers.postgres import PostgresContainer
|
||||
from testcontainers.redis import RedisContainer
|
||||
|
||||
from app_factory import create_app
|
||||
from models import db
|
||||
|
||||
# Configure logging for test containers
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DifyTestContainers:
|
||||
"""
|
||||
Manages all test containers required for Dify integration tests.
|
||||
|
||||
This class provides a centralized way to manage multiple containers
|
||||
needed for comprehensive integration testing, including databases,
|
||||
caches, and search engines.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize container management with default configurations."""
|
||||
self.postgres: Optional[PostgresContainer] = None
|
||||
self.redis: Optional[RedisContainer] = None
|
||||
self.dify_sandbox: Optional[DockerContainer] = None
|
||||
self._containers_started = False
|
||||
logger.info("DifyTestContainers initialized - ready to manage test containers")
|
||||
|
||||
def start_containers_with_env(self) -> None:
|
||||
"""
|
||||
Start all required containers for integration testing.
|
||||
|
||||
This method initializes and starts PostgreSQL, Redis
|
||||
containers with appropriate configurations for Dify testing. Containers
|
||||
are started in dependency order to ensure proper initialization.
|
||||
"""
|
||||
if self._containers_started:
|
||||
logger.info("Containers already started - skipping container startup")
|
||||
return
|
||||
|
||||
logger.info("Starting test containers for Dify integration tests...")
|
||||
|
||||
# Start PostgreSQL container for main application database
|
||||
# PostgreSQL is used for storing user data, workflows, and application state
|
||||
logger.info("Initializing PostgreSQL container...")
|
||||
self.postgres = PostgresContainer(
|
||||
image="postgres:16-alpine",
|
||||
)
|
||||
self.postgres.start()
|
||||
db_host = self.postgres.get_container_host_ip()
|
||||
db_port = self.postgres.get_exposed_port(5432)
|
||||
os.environ["DB_HOST"] = db_host
|
||||
os.environ["DB_PORT"] = str(db_port)
|
||||
os.environ["DB_USERNAME"] = self.postgres.username
|
||||
os.environ["DB_PASSWORD"] = self.postgres.password
|
||||
os.environ["DB_DATABASE"] = self.postgres.dbname
|
||||
logger.info(
|
||||
"PostgreSQL container started successfully - Host: %s, Port: %s User: %s, Database: %s",
|
||||
db_host,
|
||||
db_port,
|
||||
self.postgres.username,
|
||||
self.postgres.dbname,
|
||||
)
|
||||
|
||||
# Wait for PostgreSQL to be ready
|
||||
logger.info("Waiting for PostgreSQL to be ready to accept connections...")
|
||||
wait_for_logs(self.postgres, "is ready to accept connections", timeout=30)
|
||||
logger.info("PostgreSQL container is ready and accepting connections")
|
||||
|
||||
# Install uuid-ossp extension for UUID generation
|
||||
logger.info("Installing uuid-ossp extension...")
|
||||
try:
|
||||
import psycopg2
|
||||
|
||||
conn = psycopg2.connect(
|
||||
host=db_host,
|
||||
port=db_port,
|
||||
user=self.postgres.username,
|
||||
password=self.postgres.password,
|
||||
database=self.postgres.dbname,
|
||||
)
|
||||
conn.autocommit = True
|
||||
cursor = conn.cursor()
|
||||
cursor.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp";')
|
||||
cursor.close()
|
||||
conn.close()
|
||||
logger.info("uuid-ossp extension installed successfully")
|
||||
except Exception as e:
|
||||
logger.warning("Failed to install uuid-ossp extension: %s", e)
|
||||
|
||||
# Set up storage environment variables
|
||||
os.environ["STORAGE_TYPE"] = "opendal"
|
||||
os.environ["OPENDAL_SCHEME"] = "fs"
|
||||
os.environ["OPENDAL_FS_ROOT"] = "storage"
|
||||
|
||||
# Start Redis container for caching and session management
|
||||
# Redis is used for storing session data, cache entries, and temporary data
|
||||
logger.info("Initializing Redis container...")
|
||||
self.redis = RedisContainer(image="redis:latest", port=6379)
|
||||
self.redis.start()
|
||||
redis_host = self.redis.get_container_host_ip()
|
||||
redis_port = self.redis.get_exposed_port(6379)
|
||||
os.environ["REDIS_HOST"] = redis_host
|
||||
os.environ["REDIS_PORT"] = str(redis_port)
|
||||
logger.info("Redis container started successfully - Host: %s, Port: %s", redis_host, redis_port)
|
||||
|
||||
# Wait for Redis to be ready
|
||||
logger.info("Waiting for Redis to be ready to accept connections...")
|
||||
wait_for_logs(self.redis, "Ready to accept connections", timeout=30)
|
||||
logger.info("Redis container is ready and accepting connections")
|
||||
|
||||
# Start Dify Sandbox container for code execution environment
|
||||
# Dify Sandbox provides a secure environment for executing user code
|
||||
logger.info("Initializing Dify Sandbox container...")
|
||||
self.dify_sandbox = DockerContainer(image="langgenius/dify-sandbox:latest")
|
||||
self.dify_sandbox.with_exposed_ports(8194)
|
||||
self.dify_sandbox.env = {
|
||||
"API_KEY": "test_api_key",
|
||||
}
|
||||
self.dify_sandbox.start()
|
||||
sandbox_host = self.dify_sandbox.get_container_host_ip()
|
||||
sandbox_port = self.dify_sandbox.get_exposed_port(8194)
|
||||
os.environ["CODE_EXECUTION_ENDPOINT"] = f"http://{sandbox_host}:{sandbox_port}"
|
||||
os.environ["CODE_EXECUTION_API_KEY"] = "test_api_key"
|
||||
logger.info("Dify Sandbox container started successfully - Host: %s, Port: %s", sandbox_host, sandbox_port)
|
||||
|
||||
# Wait for Dify Sandbox to be ready
|
||||
logger.info("Waiting for Dify Sandbox to be ready to accept connections...")
|
||||
wait_for_logs(self.dify_sandbox, "config init success", timeout=60)
|
||||
logger.info("Dify Sandbox container is ready and accepting connections")
|
||||
|
||||
self._containers_started = True
|
||||
logger.info("All test containers started successfully")
|
||||
|
||||
def stop_containers(self) -> None:
|
||||
"""
|
||||
Stop and clean up all test containers.
|
||||
|
||||
This method ensures proper cleanup of all containers to prevent
|
||||
resource leaks and conflicts between test runs.
|
||||
"""
|
||||
if not self._containers_started:
|
||||
logger.info("No containers to stop - containers were not started")
|
||||
return
|
||||
|
||||
logger.info("Stopping and cleaning up test containers...")
|
||||
containers = [self.redis, self.postgres, self.dify_sandbox]
|
||||
for container in containers:
|
||||
if container:
|
||||
try:
|
||||
container_name = container.image
|
||||
logger.info("Stopping container: %s", container_name)
|
||||
container.stop()
|
||||
logger.info("Successfully stopped container: %s", container_name)
|
||||
except Exception as e:
|
||||
# Log error but don't fail the test cleanup
|
||||
logger.warning("Failed to stop container %s: %s", container, e)
|
||||
|
||||
self._containers_started = False
|
||||
logger.info("All test containers stopped and cleaned up successfully")
|
||||
|
||||
|
||||
# Global container manager instance
|
||||
_container_manager = DifyTestContainers()
|
||||
|
||||
|
||||
def _create_app_with_containers() -> Flask:
|
||||
"""
|
||||
Create Flask application configured to use test containers.
|
||||
|
||||
This function creates a Flask application instance that is configured
|
||||
to connect to the test containers instead of the default development
|
||||
or production databases.
|
||||
|
||||
Returns:
|
||||
Flask: Configured Flask application for containerized testing
|
||||
"""
|
||||
logger.info("Creating Flask application with test container configuration...")
|
||||
|
||||
# Re-create the config after environment variables have been set
|
||||
from configs import dify_config
|
||||
|
||||
# Force re-creation of config with new environment variables
|
||||
dify_config.__dict__.clear()
|
||||
dify_config.__init__()
|
||||
|
||||
# Create and configure the Flask application
|
||||
logger.info("Initializing Flask application...")
|
||||
app = create_app()
|
||||
logger.info("Flask application created successfully")
|
||||
|
||||
# Initialize database schema
|
||||
logger.info("Creating database schema...")
|
||||
with app.app_context():
|
||||
db.create_all()
|
||||
logger.info("Database schema created successfully")
|
||||
|
||||
logger.info("Flask application configured and ready for testing")
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def set_up_containers_and_env() -> Generator[DifyTestContainers, None, None]:
|
||||
"""
|
||||
Session-scoped fixture to manage test containers.
|
||||
|
||||
This fixture ensures containers are started once per test session
|
||||
and properly cleaned up when all tests are complete. This approach
|
||||
improves test performance by reusing containers across multiple tests.
|
||||
|
||||
Yields:
|
||||
DifyTestContainers: Container manager instance
|
||||
"""
|
||||
logger.info("=== Starting test session container management ===")
|
||||
_container_manager.start_containers_with_env()
|
||||
logger.info("Test containers ready for session")
|
||||
yield _container_manager
|
||||
logger.info("=== Cleaning up test session containers ===")
|
||||
_container_manager.stop_containers()
|
||||
logger.info("Test session container cleanup completed")
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def flask_app_with_containers(set_up_containers_and_env) -> Flask:
|
||||
"""
|
||||
Session-scoped Flask application fixture using test containers.
|
||||
|
||||
This fixture provides a Flask application instance that is configured
|
||||
to use the test containers for all database and service connections.
|
||||
|
||||
Args:
|
||||
containers: Container manager fixture
|
||||
|
||||
Returns:
|
||||
Flask: Configured Flask application
|
||||
"""
|
||||
logger.info("=== Creating session-scoped Flask application ===")
|
||||
app = _create_app_with_containers()
|
||||
logger.info("Session-scoped Flask application created successfully")
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def flask_req_ctx_with_containers(flask_app_with_containers) -> Generator[None, None, None]:
|
||||
"""
|
||||
Request context fixture for containerized Flask application.
|
||||
|
||||
This fixture provides a Flask request context for tests that need
|
||||
to interact with the Flask application within a request scope.
|
||||
|
||||
Args:
|
||||
flask_app_with_containers: Flask application fixture
|
||||
|
||||
Yields:
|
||||
None: Request context is active during yield
|
||||
"""
|
||||
logger.debug("Creating Flask request context...")
|
||||
with flask_app_with_containers.test_request_context():
|
||||
logger.debug("Flask request context active")
|
||||
yield
|
||||
logger.debug("Flask request context closed")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_client_with_containers(flask_app_with_containers) -> Generator[FlaskClient, None, None]:
|
||||
"""
|
||||
Test client fixture for containerized Flask application.
|
||||
|
||||
This fixture provides a Flask test client that can be used to make
|
||||
HTTP requests to the containerized application for integration testing.
|
||||
|
||||
Args:
|
||||
flask_app_with_containers: Flask application fixture
|
||||
|
||||
Yields:
|
||||
FlaskClient: Test client instance
|
||||
"""
|
||||
logger.debug("Creating Flask test client...")
|
||||
with flask_app_with_containers.test_client() as client:
|
||||
logger.debug("Flask test client ready")
|
||||
yield client
|
||||
logger.debug("Flask test client closed")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def db_session_with_containers(flask_app_with_containers) -> Generator[Session, None, None]:
|
||||
"""
|
||||
Database session fixture for containerized testing.
|
||||
|
||||
This fixture provides a SQLAlchemy database session that is connected
|
||||
to the test PostgreSQL container, allowing tests to interact with
|
||||
the database directly.
|
||||
|
||||
Args:
|
||||
flask_app_with_containers: Flask application fixture
|
||||
|
||||
Yields:
|
||||
Session: Database session instance
|
||||
"""
|
||||
logger.debug("Creating database session...")
|
||||
with flask_app_with_containers.app_context():
|
||||
session = db.session()
|
||||
logger.debug("Database session created and ready")
|
||||
try:
|
||||
yield session
|
||||
finally:
|
||||
session.close()
|
||||
logger.debug("Database session closed")
|
@@ -0,0 +1,371 @@
|
||||
import unittest
|
||||
from datetime import UTC, datetime
|
||||
from typing import Optional
|
||||
from unittest.mock import patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.file import File, FileTransferMethod, FileType
|
||||
from extensions.ext_database import db
|
||||
from factories.file_factory import StorageKeyLoader
|
||||
from models import ToolFile, UploadFile
|
||||
from models.enums import CreatorUserRole
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("flask_req_ctx_with_containers")
|
||||
class TestStorageKeyLoader(unittest.TestCase):
|
||||
"""
|
||||
Integration tests for StorageKeyLoader class.
|
||||
|
||||
Tests the batched loading of storage keys from the database for files
|
||||
with different transfer methods: LOCAL_FILE, REMOTE_URL, and TOOL_FILE.
|
||||
"""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test data before each test method."""
|
||||
self.session = db.session()
|
||||
self.tenant_id = str(uuid4())
|
||||
self.user_id = str(uuid4())
|
||||
self.conversation_id = str(uuid4())
|
||||
|
||||
# Create test data that will be cleaned up after each test
|
||||
self.test_upload_files = []
|
||||
self.test_tool_files = []
|
||||
|
||||
# Create StorageKeyLoader instance
|
||||
self.loader = StorageKeyLoader(self.session, self.tenant_id)
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up test data after each test method."""
|
||||
self.session.rollback()
|
||||
|
||||
def _create_upload_file(
|
||||
self, file_id: Optional[str] = None, storage_key: Optional[str] = None, tenant_id: Optional[str] = None
|
||||
) -> UploadFile:
|
||||
"""Helper method to create an UploadFile record for testing."""
|
||||
if file_id is None:
|
||||
file_id = str(uuid4())
|
||||
if storage_key is None:
|
||||
storage_key = f"test_storage_key_{uuid4()}"
|
||||
if tenant_id is None:
|
||||
tenant_id = self.tenant_id
|
||||
|
||||
upload_file = UploadFile(
|
||||
tenant_id=tenant_id,
|
||||
storage_type="local",
|
||||
key=storage_key,
|
||||
name="test_file.txt",
|
||||
size=1024,
|
||||
extension=".txt",
|
||||
mime_type="text/plain",
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=self.user_id,
|
||||
created_at=datetime.now(UTC),
|
||||
used=False,
|
||||
)
|
||||
upload_file.id = file_id
|
||||
|
||||
self.session.add(upload_file)
|
||||
self.session.flush()
|
||||
self.test_upload_files.append(upload_file)
|
||||
|
||||
return upload_file
|
||||
|
||||
def _create_tool_file(
|
||||
self, file_id: Optional[str] = None, file_key: Optional[str] = None, tenant_id: Optional[str] = None
|
||||
) -> ToolFile:
|
||||
"""Helper method to create a ToolFile record for testing."""
|
||||
if file_id is None:
|
||||
file_id = str(uuid4())
|
||||
if file_key is None:
|
||||
file_key = f"test_file_key_{uuid4()}"
|
||||
if tenant_id is None:
|
||||
tenant_id = self.tenant_id
|
||||
|
||||
tool_file = ToolFile()
|
||||
tool_file.id = file_id
|
||||
tool_file.user_id = self.user_id
|
||||
tool_file.tenant_id = tenant_id
|
||||
tool_file.conversation_id = self.conversation_id
|
||||
tool_file.file_key = file_key
|
||||
tool_file.mimetype = "text/plain"
|
||||
tool_file.original_url = "http://example.com/file.txt"
|
||||
tool_file.name = "test_tool_file.txt"
|
||||
tool_file.size = 2048
|
||||
|
||||
self.session.add(tool_file)
|
||||
self.session.flush()
|
||||
self.test_tool_files.append(tool_file)
|
||||
|
||||
return tool_file
|
||||
|
||||
def _create_file(
|
||||
self, related_id: str, transfer_method: FileTransferMethod, tenant_id: Optional[str] = None
|
||||
) -> File:
|
||||
"""Helper method to create a File object for testing."""
|
||||
if tenant_id is None:
|
||||
tenant_id = self.tenant_id
|
||||
|
||||
# Set related_id for LOCAL_FILE and TOOL_FILE transfer methods
|
||||
file_related_id = None
|
||||
remote_url = None
|
||||
|
||||
if transfer_method in (FileTransferMethod.LOCAL_FILE, FileTransferMethod.TOOL_FILE):
|
||||
file_related_id = related_id
|
||||
elif transfer_method == FileTransferMethod.REMOTE_URL:
|
||||
remote_url = "https://example.com/test_file.txt"
|
||||
file_related_id = related_id
|
||||
|
||||
return File(
|
||||
id=str(uuid4()), # Generate new UUID for File.id
|
||||
tenant_id=tenant_id,
|
||||
type=FileType.DOCUMENT,
|
||||
transfer_method=transfer_method,
|
||||
related_id=file_related_id,
|
||||
remote_url=remote_url,
|
||||
filename="test_file.txt",
|
||||
extension=".txt",
|
||||
mime_type="text/plain",
|
||||
size=1024,
|
||||
storage_key="initial_key",
|
||||
)
|
||||
|
||||
def test_load_storage_keys_local_file(self):
|
||||
"""Test loading storage keys for LOCAL_FILE transfer method."""
|
||||
# Create test data
|
||||
upload_file = self._create_upload_file()
|
||||
file = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE)
|
||||
|
||||
# Load storage keys
|
||||
self.loader.load_storage_keys([file])
|
||||
|
||||
# Verify storage key was loaded correctly
|
||||
assert file._storage_key == upload_file.key
|
||||
|
||||
def test_load_storage_keys_remote_url(self):
|
||||
"""Test loading storage keys for REMOTE_URL transfer method."""
|
||||
# Create test data
|
||||
upload_file = self._create_upload_file()
|
||||
file = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.REMOTE_URL)
|
||||
|
||||
# Load storage keys
|
||||
self.loader.load_storage_keys([file])
|
||||
|
||||
# Verify storage key was loaded correctly
|
||||
assert file._storage_key == upload_file.key
|
||||
|
||||
def test_load_storage_keys_tool_file(self):
|
||||
"""Test loading storage keys for TOOL_FILE transfer method."""
|
||||
# Create test data
|
||||
tool_file = self._create_tool_file()
|
||||
file = self._create_file(related_id=tool_file.id, transfer_method=FileTransferMethod.TOOL_FILE)
|
||||
|
||||
# Load storage keys
|
||||
self.loader.load_storage_keys([file])
|
||||
|
||||
# Verify storage key was loaded correctly
|
||||
assert file._storage_key == tool_file.file_key
|
||||
|
||||
def test_load_storage_keys_mixed_methods(self):
|
||||
"""Test batch loading with mixed transfer methods."""
|
||||
# Create test data for different transfer methods
|
||||
upload_file1 = self._create_upload_file()
|
||||
upload_file2 = self._create_upload_file()
|
||||
tool_file = self._create_tool_file()
|
||||
|
||||
file1 = self._create_file(related_id=upload_file1.id, transfer_method=FileTransferMethod.LOCAL_FILE)
|
||||
file2 = self._create_file(related_id=upload_file2.id, transfer_method=FileTransferMethod.REMOTE_URL)
|
||||
file3 = self._create_file(related_id=tool_file.id, transfer_method=FileTransferMethod.TOOL_FILE)
|
||||
|
||||
files = [file1, file2, file3]
|
||||
|
||||
# Load storage keys
|
||||
self.loader.load_storage_keys(files)
|
||||
|
||||
# Verify all storage keys were loaded correctly
|
||||
assert file1._storage_key == upload_file1.key
|
||||
assert file2._storage_key == upload_file2.key
|
||||
assert file3._storage_key == tool_file.file_key
|
||||
|
||||
def test_load_storage_keys_empty_list(self):
|
||||
"""Test with empty file list."""
|
||||
# Should not raise any exceptions
|
||||
self.loader.load_storage_keys([])
|
||||
|
||||
def test_load_storage_keys_tenant_mismatch(self):
|
||||
"""Test tenant_id validation."""
|
||||
# Create file with different tenant_id
|
||||
upload_file = self._create_upload_file()
|
||||
file = self._create_file(
|
||||
related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE, tenant_id=str(uuid4())
|
||||
)
|
||||
|
||||
# Should raise ValueError for tenant mismatch
|
||||
with pytest.raises(ValueError) as context:
|
||||
self.loader.load_storage_keys([file])
|
||||
|
||||
assert "invalid file, expected tenant_id" in str(context.value)
|
||||
|
||||
def test_load_storage_keys_missing_file_id(self):
|
||||
"""Test with None file.related_id."""
|
||||
# Create a file with valid parameters first, then manually set related_id to None
|
||||
file = self._create_file(related_id=str(uuid4()), transfer_method=FileTransferMethod.LOCAL_FILE)
|
||||
file.related_id = None
|
||||
|
||||
# Should raise ValueError for None file related_id
|
||||
with pytest.raises(ValueError) as context:
|
||||
self.loader.load_storage_keys([file])
|
||||
|
||||
assert str(context.value) == "file id should not be None."
|
||||
|
||||
def test_load_storage_keys_nonexistent_upload_file_records(self):
|
||||
"""Test with missing UploadFile database records."""
|
||||
# Create file with non-existent upload file id
|
||||
non_existent_id = str(uuid4())
|
||||
file = self._create_file(related_id=non_existent_id, transfer_method=FileTransferMethod.LOCAL_FILE)
|
||||
|
||||
# Should raise ValueError for missing record
|
||||
with pytest.raises(ValueError):
|
||||
self.loader.load_storage_keys([file])
|
||||
|
||||
def test_load_storage_keys_nonexistent_tool_file_records(self):
|
||||
"""Test with missing ToolFile database records."""
|
||||
# Create file with non-existent tool file id
|
||||
non_existent_id = str(uuid4())
|
||||
file = self._create_file(related_id=non_existent_id, transfer_method=FileTransferMethod.TOOL_FILE)
|
||||
|
||||
# Should raise ValueError for missing record
|
||||
with pytest.raises(ValueError):
|
||||
self.loader.load_storage_keys([file])
|
||||
|
||||
def test_load_storage_keys_invalid_uuid(self):
|
||||
"""Test with invalid UUID format."""
|
||||
# Create a file with valid parameters first, then manually set invalid related_id
|
||||
file = self._create_file(related_id=str(uuid4()), transfer_method=FileTransferMethod.LOCAL_FILE)
|
||||
file.related_id = "invalid-uuid-format"
|
||||
|
||||
# Should raise ValueError for invalid UUID
|
||||
with pytest.raises(ValueError):
|
||||
self.loader.load_storage_keys([file])
|
||||
|
||||
def test_load_storage_keys_batch_efficiency(self):
|
||||
"""Test batched operations use efficient queries."""
|
||||
# Create multiple files of different types
|
||||
upload_files = [self._create_upload_file() for _ in range(3)]
|
||||
tool_files = [self._create_tool_file() for _ in range(2)]
|
||||
|
||||
files = []
|
||||
files.extend(
|
||||
[self._create_file(related_id=uf.id, transfer_method=FileTransferMethod.LOCAL_FILE) for uf in upload_files]
|
||||
)
|
||||
files.extend(
|
||||
[self._create_file(related_id=tf.id, transfer_method=FileTransferMethod.TOOL_FILE) for tf in tool_files]
|
||||
)
|
||||
|
||||
# Mock the session to count queries
|
||||
with patch.object(self.session, "scalars", wraps=self.session.scalars) as mock_scalars:
|
||||
self.loader.load_storage_keys(files)
|
||||
|
||||
# Should make exactly 2 queries (one for upload_files, one for tool_files)
|
||||
assert mock_scalars.call_count == 2
|
||||
|
||||
# Verify all storage keys were loaded correctly
|
||||
for i, file in enumerate(files[:3]):
|
||||
assert file._storage_key == upload_files[i].key
|
||||
for i, file in enumerate(files[3:]):
|
||||
assert file._storage_key == tool_files[i].file_key
|
||||
|
||||
def test_load_storage_keys_tenant_isolation(self):
|
||||
"""Test that tenant isolation works correctly."""
|
||||
# Create files for different tenants
|
||||
other_tenant_id = str(uuid4())
|
||||
|
||||
# Create upload file for current tenant
|
||||
upload_file_current = self._create_upload_file()
|
||||
file_current = self._create_file(
|
||||
related_id=upload_file_current.id, transfer_method=FileTransferMethod.LOCAL_FILE
|
||||
)
|
||||
|
||||
# Create upload file for other tenant (but don't add to cleanup list)
|
||||
upload_file_other = UploadFile(
|
||||
tenant_id=other_tenant_id,
|
||||
storage_type="local",
|
||||
key="other_tenant_key",
|
||||
name="other_file.txt",
|
||||
size=1024,
|
||||
extension=".txt",
|
||||
mime_type="text/plain",
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by=self.user_id,
|
||||
created_at=datetime.now(UTC),
|
||||
used=False,
|
||||
)
|
||||
upload_file_other.id = str(uuid4())
|
||||
self.session.add(upload_file_other)
|
||||
self.session.flush()
|
||||
|
||||
# Create file for other tenant but try to load with current tenant's loader
|
||||
file_other = self._create_file(
|
||||
related_id=upload_file_other.id, transfer_method=FileTransferMethod.LOCAL_FILE, tenant_id=other_tenant_id
|
||||
)
|
||||
|
||||
# Should raise ValueError due to tenant mismatch
|
||||
with pytest.raises(ValueError) as context:
|
||||
self.loader.load_storage_keys([file_other])
|
||||
|
||||
assert "invalid file, expected tenant_id" in str(context.value)
|
||||
|
||||
# Current tenant's file should still work
|
||||
self.loader.load_storage_keys([file_current])
|
||||
assert file_current._storage_key == upload_file_current.key
|
||||
|
||||
def test_load_storage_keys_mixed_tenant_batch(self):
|
||||
"""Test batch with mixed tenant files (should fail on first mismatch)."""
|
||||
# Create files for current tenant
|
||||
upload_file_current = self._create_upload_file()
|
||||
file_current = self._create_file(
|
||||
related_id=upload_file_current.id, transfer_method=FileTransferMethod.LOCAL_FILE
|
||||
)
|
||||
|
||||
# Create file for different tenant
|
||||
other_tenant_id = str(uuid4())
|
||||
file_other = self._create_file(
|
||||
related_id=str(uuid4()), transfer_method=FileTransferMethod.LOCAL_FILE, tenant_id=other_tenant_id
|
||||
)
|
||||
|
||||
# Should raise ValueError on tenant mismatch
|
||||
with pytest.raises(ValueError) as context:
|
||||
self.loader.load_storage_keys([file_current, file_other])
|
||||
|
||||
assert "invalid file, expected tenant_id" in str(context.value)
|
||||
|
||||
def test_load_storage_keys_duplicate_file_ids(self):
|
||||
"""Test handling of duplicate file IDs in the batch."""
|
||||
# Create upload file
|
||||
upload_file = self._create_upload_file()
|
||||
|
||||
# Create two File objects with same related_id
|
||||
file1 = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE)
|
||||
file2 = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE)
|
||||
|
||||
# Should handle duplicates gracefully
|
||||
self.loader.load_storage_keys([file1, file2])
|
||||
|
||||
# Both files should have the same storage key
|
||||
assert file1._storage_key == upload_file.key
|
||||
assert file2._storage_key == upload_file.key
|
||||
|
||||
def test_load_storage_keys_session_isolation(self):
|
||||
"""Test that the loader uses the provided session correctly."""
|
||||
# Create test data
|
||||
upload_file = self._create_upload_file()
|
||||
file = self._create_file(related_id=upload_file.id, transfer_method=FileTransferMethod.LOCAL_FILE)
|
||||
|
||||
# Create loader with different session (same underlying connection)
|
||||
|
||||
with Session(bind=db.engine) as other_session:
|
||||
other_loader = StorageKeyLoader(other_session, self.tenant_id)
|
||||
with pytest.raises(ValueError):
|
||||
other_loader.load_storage_keys([file])
|
@@ -0,0 +1,11 @@
|
||||
import pytest
|
||||
|
||||
from core.helper.code_executor.code_executor import CodeExecutionError, CodeExecutor
|
||||
|
||||
CODE_LANGUAGE = "unsupported_language"
|
||||
|
||||
|
||||
def test_unsupported_with_code_template():
|
||||
with pytest.raises(CodeExecutionError) as e:
|
||||
CodeExecutor.execute_workflow_code_template(language=CODE_LANGUAGE, code="", inputs={})
|
||||
assert str(e.value) == f"Unsupported language {CODE_LANGUAGE}"
|
@@ -0,0 +1,47 @@
|
||||
from textwrap import dedent
|
||||
|
||||
from .test_utils import CodeExecutorTestMixin
|
||||
|
||||
|
||||
class TestJavaScriptCodeExecutor(CodeExecutorTestMixin):
|
||||
"""Test class for JavaScript code executor functionality."""
|
||||
|
||||
def test_javascript_plain(self, flask_app_with_containers):
|
||||
"""Test basic JavaScript code execution with console.log output"""
|
||||
CodeExecutor, CodeLanguage = self.code_executor_imports
|
||||
|
||||
code = 'console.log("Hello World")'
|
||||
result_message = CodeExecutor.execute_code(language=CodeLanguage.JAVASCRIPT, preload="", code=code)
|
||||
assert result_message == "Hello World\n"
|
||||
|
||||
def test_javascript_json(self, flask_app_with_containers):
|
||||
"""Test JavaScript code execution with JSON output"""
|
||||
CodeExecutor, CodeLanguage = self.code_executor_imports
|
||||
|
||||
code = dedent("""
|
||||
obj = {'Hello': 'World'}
|
||||
console.log(JSON.stringify(obj))
|
||||
""")
|
||||
result = CodeExecutor.execute_code(language=CodeLanguage.JAVASCRIPT, preload="", code=code)
|
||||
assert result == '{"Hello":"World"}\n'
|
||||
|
||||
def test_javascript_with_code_template(self, flask_app_with_containers):
|
||||
"""Test JavaScript workflow code template execution with inputs"""
|
||||
CodeExecutor, CodeLanguage = self.code_executor_imports
|
||||
JavascriptCodeProvider, _ = self.javascript_imports
|
||||
|
||||
result = CodeExecutor.execute_workflow_code_template(
|
||||
language=CodeLanguage.JAVASCRIPT,
|
||||
code=JavascriptCodeProvider.get_default_code(),
|
||||
inputs={"arg1": "Hello", "arg2": "World"},
|
||||
)
|
||||
assert result == {"result": "HelloWorld"}
|
||||
|
||||
def test_javascript_get_runner_script(self, flask_app_with_containers):
|
||||
"""Test JavaScript template transformer runner script generation"""
|
||||
_, NodeJsTemplateTransformer = self.javascript_imports
|
||||
|
||||
runner_script = NodeJsTemplateTransformer.get_runner_script()
|
||||
assert runner_script.count(NodeJsTemplateTransformer._code_placeholder) == 1
|
||||
assert runner_script.count(NodeJsTemplateTransformer._inputs_placeholder) == 1
|
||||
assert runner_script.count(NodeJsTemplateTransformer._result_tag) == 2
|
@@ -0,0 +1,42 @@
|
||||
import base64
|
||||
|
||||
from .test_utils import CodeExecutorTestMixin
|
||||
|
||||
|
||||
class TestJinja2CodeExecutor(CodeExecutorTestMixin):
|
||||
"""Test class for Jinja2 code executor functionality."""
|
||||
|
||||
def test_jinja2(self, flask_app_with_containers):
|
||||
"""Test basic Jinja2 template execution with variable substitution"""
|
||||
CodeExecutor, CodeLanguage = self.code_executor_imports
|
||||
_, Jinja2TemplateTransformer = self.jinja2_imports
|
||||
|
||||
template = "Hello {{template}}"
|
||||
inputs = base64.b64encode(b'{"template": "World"}').decode("utf-8")
|
||||
code = (
|
||||
Jinja2TemplateTransformer.get_runner_script()
|
||||
.replace(Jinja2TemplateTransformer._code_placeholder, template)
|
||||
.replace(Jinja2TemplateTransformer._inputs_placeholder, inputs)
|
||||
)
|
||||
result = CodeExecutor.execute_code(
|
||||
language=CodeLanguage.JINJA2, preload=Jinja2TemplateTransformer.get_preload_script(), code=code
|
||||
)
|
||||
assert result == "<<RESULT>>Hello World<<RESULT>>\n"
|
||||
|
||||
def test_jinja2_with_code_template(self, flask_app_with_containers):
|
||||
"""Test Jinja2 workflow code template execution with inputs"""
|
||||
CodeExecutor, CodeLanguage = self.code_executor_imports
|
||||
|
||||
result = CodeExecutor.execute_workflow_code_template(
|
||||
language=CodeLanguage.JINJA2, code="Hello {{template}}", inputs={"template": "World"}
|
||||
)
|
||||
assert result == {"result": "Hello World"}
|
||||
|
||||
def test_jinja2_get_runner_script(self, flask_app_with_containers):
|
||||
"""Test Jinja2 template transformer runner script generation"""
|
||||
_, Jinja2TemplateTransformer = self.jinja2_imports
|
||||
|
||||
runner_script = Jinja2TemplateTransformer.get_runner_script()
|
||||
assert runner_script.count(Jinja2TemplateTransformer._code_placeholder) == 1
|
||||
assert runner_script.count(Jinja2TemplateTransformer._inputs_placeholder) == 1
|
||||
assert runner_script.count(Jinja2TemplateTransformer._result_tag) == 2
|
@@ -0,0 +1,47 @@
|
||||
from textwrap import dedent
|
||||
|
||||
from .test_utils import CodeExecutorTestMixin
|
||||
|
||||
|
||||
class TestPython3CodeExecutor(CodeExecutorTestMixin):
|
||||
"""Test class for Python3 code executor functionality."""
|
||||
|
||||
def test_python3_plain(self, flask_app_with_containers):
|
||||
"""Test basic Python3 code execution with print output"""
|
||||
CodeExecutor, CodeLanguage = self.code_executor_imports
|
||||
|
||||
code = 'print("Hello World")'
|
||||
result = CodeExecutor.execute_code(language=CodeLanguage.PYTHON3, preload="", code=code)
|
||||
assert result == "Hello World\n"
|
||||
|
||||
def test_python3_json(self, flask_app_with_containers):
|
||||
"""Test Python3 code execution with JSON output"""
|
||||
CodeExecutor, CodeLanguage = self.code_executor_imports
|
||||
|
||||
code = dedent("""
|
||||
import json
|
||||
print(json.dumps({'Hello': 'World'}))
|
||||
""")
|
||||
result = CodeExecutor.execute_code(language=CodeLanguage.PYTHON3, preload="", code=code)
|
||||
assert result == '{"Hello": "World"}\n'
|
||||
|
||||
def test_python3_with_code_template(self, flask_app_with_containers):
|
||||
"""Test Python3 workflow code template execution with inputs"""
|
||||
CodeExecutor, CodeLanguage = self.code_executor_imports
|
||||
Python3CodeProvider, _ = self.python3_imports
|
||||
|
||||
result = CodeExecutor.execute_workflow_code_template(
|
||||
language=CodeLanguage.PYTHON3,
|
||||
code=Python3CodeProvider.get_default_code(),
|
||||
inputs={"arg1": "Hello", "arg2": "World"},
|
||||
)
|
||||
assert result == {"result": "HelloWorld"}
|
||||
|
||||
def test_python3_get_runner_script(self, flask_app_with_containers):
|
||||
"""Test Python3 template transformer runner script generation"""
|
||||
_, Python3TemplateTransformer = self.python3_imports
|
||||
|
||||
runner_script = Python3TemplateTransformer.get_runner_script()
|
||||
assert runner_script.count(Python3TemplateTransformer._code_placeholder) == 1
|
||||
assert runner_script.count(Python3TemplateTransformer._inputs_placeholder) == 1
|
||||
assert runner_script.count(Python3TemplateTransformer._result_tag) == 2
|
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
Test utilities for code executor integration tests.
|
||||
|
||||
This module provides lazy import functions to avoid module loading issues
|
||||
that occur when modules are imported before the flask_app_with_containers fixture
|
||||
has set up the proper environment variables and configuration.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
|
||||
def force_reload_code_executor():
|
||||
"""
|
||||
Force reload the code_executor module to reinitialize code_execution_endpoint_url.
|
||||
|
||||
This function should be called after setting up environment variables
|
||||
to ensure the code_execution_endpoint_url is initialized with the correct value.
|
||||
"""
|
||||
try:
|
||||
import core.helper.code_executor.code_executor
|
||||
|
||||
importlib.reload(core.helper.code_executor.code_executor)
|
||||
except Exception as e:
|
||||
# Log the error but don't fail the test
|
||||
print(f"Warning: Failed to reload code_executor module: {e}")
|
||||
|
||||
|
||||
def get_code_executor_imports():
|
||||
"""
|
||||
Lazy import function for core CodeExecutor classes.
|
||||
|
||||
Returns:
|
||||
tuple: (CodeExecutor, CodeLanguage) classes
|
||||
"""
|
||||
from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
|
||||
|
||||
return CodeExecutor, CodeLanguage
|
||||
|
||||
|
||||
def get_javascript_imports():
|
||||
"""
|
||||
Lazy import function for JavaScript-specific modules.
|
||||
|
||||
Returns:
|
||||
tuple: (JavascriptCodeProvider, NodeJsTemplateTransformer) classes
|
||||
"""
|
||||
from core.helper.code_executor.javascript.javascript_code_provider import JavascriptCodeProvider
|
||||
from core.helper.code_executor.javascript.javascript_transformer import NodeJsTemplateTransformer
|
||||
|
||||
return JavascriptCodeProvider, NodeJsTemplateTransformer
|
||||
|
||||
|
||||
def get_python3_imports():
|
||||
"""
|
||||
Lazy import function for Python3-specific modules.
|
||||
|
||||
Returns:
|
||||
tuple: (Python3CodeProvider, Python3TemplateTransformer) classes
|
||||
"""
|
||||
from core.helper.code_executor.python3.python3_code_provider import Python3CodeProvider
|
||||
from core.helper.code_executor.python3.python3_transformer import Python3TemplateTransformer
|
||||
|
||||
return Python3CodeProvider, Python3TemplateTransformer
|
||||
|
||||
|
||||
def get_jinja2_imports():
|
||||
"""
|
||||
Lazy import function for Jinja2-specific modules.
|
||||
|
||||
Returns:
|
||||
tuple: (None, Jinja2TemplateTransformer) classes
|
||||
"""
|
||||
from core.helper.code_executor.jinja2.jinja2_transformer import Jinja2TemplateTransformer
|
||||
|
||||
return None, Jinja2TemplateTransformer
|
||||
|
||||
|
||||
class CodeExecutorTestMixin:
|
||||
"""
|
||||
Mixin class providing lazy import methods for code executor tests.
|
||||
|
||||
This mixin helps avoid module loading issues by deferring imports
|
||||
until after the flask_app_with_containers fixture has set up the environment.
|
||||
"""
|
||||
|
||||
def setup_method(self):
|
||||
"""
|
||||
Setup method called before each test method.
|
||||
Force reload the code_executor module to ensure fresh initialization.
|
||||
"""
|
||||
force_reload_code_executor()
|
||||
|
||||
@property
|
||||
def code_executor_imports(self):
|
||||
"""Property to get CodeExecutor and CodeLanguage classes."""
|
||||
return get_code_executor_imports()
|
||||
|
||||
@property
|
||||
def javascript_imports(self):
|
||||
"""Property to get JavaScript-specific classes."""
|
||||
return get_javascript_imports()
|
||||
|
||||
@property
|
||||
def python3_imports(self):
|
||||
"""Property to get Python3-specific classes."""
|
||||
return get_python3_imports()
|
||||
|
||||
@property
|
||||
def jinja2_imports(self):
|
||||
"""Property to get Jinja2-specific classes."""
|
||||
return get_jinja2_imports()
|
32
api/uv.lock
generated
32
api/uv.lock
generated
@@ -1318,6 +1318,7 @@ dev = [
|
||||
{ name = "pytest-mock" },
|
||||
{ name = "ruff" },
|
||||
{ name = "scipy-stubs" },
|
||||
{ name = "testcontainers" },
|
||||
{ name = "types-aiofiles" },
|
||||
{ name = "types-beautifulsoup4" },
|
||||
{ name = "types-cachetools" },
|
||||
@@ -1500,6 +1501,7 @@ dev = [
|
||||
{ name = "pytest-mock", specifier = "~=3.14.0" },
|
||||
{ name = "ruff", specifier = "~=0.12.3" },
|
||||
{ name = "scipy-stubs", specifier = ">=1.15.3.0" },
|
||||
{ name = "testcontainers", specifier = "~=4.10.0" },
|
||||
{ name = "types-aiofiles", specifier = "~=24.1.0" },
|
||||
{ name = "types-beautifulsoup4", specifier = "~=4.12.0" },
|
||||
{ name = "types-cachetools", specifier = "~=5.5.0" },
|
||||
@@ -1600,6 +1602,20 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/12/b3/231ffd4ab1fc9d679809f356cebee130ac7daa00d6d6f3206dd4fd137e9e/distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2", size = 20277, upload-time = "2023-12-24T09:54:30.421Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "docker"
|
||||
version = "7.1.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "pywin32", marker = "sys_platform == 'win32'" },
|
||||
{ name = "requests" },
|
||||
{ name = "urllib3" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/91/9b/4a2ea29aeba62471211598dac5d96825bb49348fa07e906ea930394a83ce/docker-7.1.0.tar.gz", hash = "sha256:ad8c70e6e3f8926cb8a92619b832b4ea5299e2831c14284663184e200546fa6c", size = 117834, upload-time = "2024-05-23T11:13:57.216Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e3/26/57c6fb270950d476074c087527a558ccb6f4436657314bfb6cdf484114c4/docker-7.1.0-py3-none-any.whl", hash = "sha256:c96b93b7f0a746f9e77d325bcfb87422a3d8bd4f03136ae8a85b37f1898d5fc0", size = 147774, upload-time = "2024-05-23T11:13:55.01Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "docstring-parser"
|
||||
version = "0.16"
|
||||
@@ -5468,6 +5484,22 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/e5/30/643397144bfbfec6f6ef821f36f33e57d35946c44a2352d3c9f0ae847619/tenacity-9.1.2-py3-none-any.whl", hash = "sha256:f77bf36710d8b73a50b2dd155c97b870017ad21afe6ab300326b0371b3b05138", size = 28248, upload-time = "2025-04-02T08:25:07.678Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "testcontainers"
|
||||
version = "4.10.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "docker" },
|
||||
{ name = "python-dotenv" },
|
||||
{ name = "typing-extensions" },
|
||||
{ name = "urllib3" },
|
||||
{ name = "wrapt" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/a1/49/9c618aff1c50121d183cdfbc3a4a5cf2727a2cde1893efe6ca55c7009196/testcontainers-4.10.0.tar.gz", hash = "sha256:03f85c3e505d8b4edeb192c72a961cebbcba0dd94344ae778b4a159cb6dcf8d3", size = 63327, upload-time = "2025-04-02T16:13:27.582Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/1c/0a/824b0c1ecf224802125279c3effff2e25ed785ed046e67da6e53d928de4c/testcontainers-4.10.0-py3-none-any.whl", hash = "sha256:31ed1a81238c7e131a2a29df6db8f23717d892b592fa5a1977fd0dcd0c23fc23", size = 107414, upload-time = "2025-04-02T16:13:25.785Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tidb-vector"
|
||||
version = "0.0.9"
|
||||
|
@@ -15,3 +15,6 @@ dev/pytest/pytest_workflow.sh
|
||||
|
||||
# Unit tests
|
||||
dev/pytest/pytest_unit_tests.sh
|
||||
|
||||
# TestContainers tests
|
||||
dev/pytest/pytest_testcontainers.sh
|
||||
|
7
dev/pytest/pytest_testcontainers.sh
Executable file
7
dev/pytest/pytest_testcontainers.sh
Executable file
@@ -0,0 +1,7 @@
|
||||
#!/bin/bash
|
||||
set -x
|
||||
|
||||
SCRIPT_DIR="$(dirname "$(realpath "$0")")"
|
||||
cd "$SCRIPT_DIR/../.."
|
||||
|
||||
pytest api/tests/test_containers_integration_tests
|
Reference in New Issue
Block a user