From 60c7663a807353d07ad24ed6759e35ea5c642dba Mon Sep 17 00:00:00 2001 From: NeatGuyCoding <15627489+NeatGuyCoding@users.noreply.github.com> Date: Mon, 4 Aug 2025 19:27:36 +0800 Subject: [PATCH] Feat add testcontainers test (#23269) --- .github/workflows/api-tests.yml | 3 + api/pyproject.toml | 1 + .../__init__.py | 0 .../conftest.py | 328 ++++++++++++++++ .../factories/__init__.py | 0 .../factories/test_storage_key_loader.py | 371 ++++++++++++++++++ .../workflow/__init__.py | 0 .../workflow/nodes/__init__.py | 0 .../workflow/nodes/code_executor/__init__.py | 0 .../nodes/code_executor/test_code_executor.py | 11 + .../code_executor/test_code_javascript.py | 47 +++ .../nodes/code_executor/test_code_jinja2.py | 42 ++ .../nodes/code_executor/test_code_python3.py | 47 +++ .../nodes/code_executor/test_utils.py | 115 ++++++ api/uv.lock | 32 ++ dev/pytest/pytest_all_tests.sh | 3 + dev/pytest/pytest_testcontainers.sh | 7 + 17 files changed, 1007 insertions(+) create mode 100644 api/tests/test_containers_integration_tests/__init__.py create mode 100644 api/tests/test_containers_integration_tests/conftest.py create mode 100644 api/tests/test_containers_integration_tests/factories/__init__.py create mode 100644 api/tests/test_containers_integration_tests/factories/test_storage_key_loader.py create mode 100644 api/tests/test_containers_integration_tests/workflow/__init__.py create mode 100644 api/tests/test_containers_integration_tests/workflow/nodes/__init__.py create mode 100644 api/tests/test_containers_integration_tests/workflow/nodes/code_executor/__init__.py create mode 100644 api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_executor.py create mode 100644 api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_javascript.py create mode 100644 api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_jinja2.py create mode 100644 api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_python3.py create mode 100644 api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_utils.py create mode 100755 dev/pytest/pytest_testcontainers.sh diff --git a/.github/workflows/api-tests.yml b/.github/workflows/api-tests.yml index a5a5071fa..9c3daddbf 100644 --- a/.github/workflows/api-tests.yml +++ b/.github/workflows/api-tests.yml @@ -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 diff --git a/api/pyproject.toml b/api/pyproject.toml index be42b509e..d8f663ef8 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -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", diff --git a/api/tests/test_containers_integration_tests/__init__.py b/api/tests/test_containers_integration_tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/api/tests/test_containers_integration_tests/conftest.py b/api/tests/test_containers_integration_tests/conftest.py new file mode 100644 index 000000000..0369a5cbd --- /dev/null +++ b/api/tests/test_containers_integration_tests/conftest.py @@ -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") diff --git a/api/tests/test_containers_integration_tests/factories/__init__.py b/api/tests/test_containers_integration_tests/factories/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/api/tests/test_containers_integration_tests/factories/test_storage_key_loader.py b/api/tests/test_containers_integration_tests/factories/test_storage_key_loader.py new file mode 100644 index 000000000..d6e14f3f5 --- /dev/null +++ b/api/tests/test_containers_integration_tests/factories/test_storage_key_loader.py @@ -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]) diff --git a/api/tests/test_containers_integration_tests/workflow/__init__.py b/api/tests/test_containers_integration_tests/workflow/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/__init__.py b/api/tests/test_containers_integration_tests/workflow/nodes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/__init__.py b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_executor.py b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_executor.py new file mode 100644 index 000000000..487178ff5 --- /dev/null +++ b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_executor.py @@ -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}" diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_javascript.py b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_javascript.py new file mode 100644 index 000000000..19a41b618 --- /dev/null +++ b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_javascript.py @@ -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 diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_jinja2.py b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_jinja2.py new file mode 100644 index 000000000..c76480117 --- /dev/null +++ b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_jinja2.py @@ -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 == "<>Hello World<>\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 diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_python3.py b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_python3.py new file mode 100644 index 000000000..6d93df247 --- /dev/null +++ b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_code_python3.py @@ -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 diff --git a/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_utils.py b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_utils.py new file mode 100644 index 000000000..35a095b04 --- /dev/null +++ b/api/tests/test_containers_integration_tests/workflow/nodes/code_executor/test_utils.py @@ -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() diff --git a/api/uv.lock b/api/uv.lock index 0bce38812..4dced728a 100644 --- a/api/uv.lock +++ b/api/uv.lock @@ -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" diff --git a/dev/pytest/pytest_all_tests.sh b/dev/pytest/pytest_all_tests.sh index 30898b4fc..9123b2f8a 100755 --- a/dev/pytest/pytest_all_tests.sh +++ b/dev/pytest/pytest_all_tests.sh @@ -15,3 +15,6 @@ dev/pytest/pytest_workflow.sh # Unit tests dev/pytest/pytest_unit_tests.sh + +# TestContainers tests +dev/pytest/pytest_testcontainers.sh diff --git a/dev/pytest/pytest_testcontainers.sh b/dev/pytest/pytest_testcontainers.sh new file mode 100755 index 000000000..e55a43613 --- /dev/null +++ b/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