| @@ -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", | |||
| @@ -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() | |||
| @@ -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 | |||
| @@ -0,0 +1,7 @@ | |||
| #!/bin/bash | |||
| set -x | |||
| SCRIPT_DIR="$(dirname "$(realpath "$0")")" | |||
| cd "$SCRIPT_DIR/../.." | |||
| pytest api/tests/test_containers_integration_tests | |||