Feat add testcontainers test (#23269)

This commit is contained in:
NeatGuyCoding
2025-08-04 19:27:36 +08:00
committed by GitHub
parent 8041808b53
commit 60c7663a80
17 changed files with 1007 additions and 0 deletions

View File

@@ -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

View File

@@ -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",

View 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")

View File

@@ -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])

View 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}"

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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
View File

@@ -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"

View File

@@ -15,3 +15,6 @@ dev/pytest/pytest_workflow.sh
# Unit tests
dev/pytest/pytest_unit_tests.sh
# TestContainers tests
dev/pytest/pytest_testcontainers.sh

View File

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