zhaoxi
d84212f780
正式发布 Pretor 平台的首个 alpha 版本。本项目旨在构建一个基于分布式架构的多智能体协同工作流水线。 核心功能实现： 1. 建立基于 BaseIndividual 的动态插件加载机制。 2. 实现三类核心 worker_individual 子个体。 3. 集成 Ray 框架支持分布式集群调度。 4. 基于 PostgreSQL 的全量持久化存储方案。 5. 提供完整的 FastAPI 后端与 React 前端交互界面。
79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
import sys
|
|
import builtins
|
|
|
|
real_import = builtins.__import__
|
|
|
|
|
|
def mock_import(name, globals=None, locals=None, fromlist=(), level=0):
|
|
if name == 'ray':
|
|
mock_ray = MagicMock()
|
|
|
|
def mock_remote(*args, **kwargs):
|
|
if len(args) == 1 and callable(args[0]):
|
|
return args[0]
|
|
def decorator(cls):
|
|
return cls
|
|
return decorator
|
|
|
|
mock_ray.remote = mock_remote
|
|
return mock_ray
|
|
return real_import(name, globals, locals, fromlist, level)
|
|
|
|
|
|
builtins.__import__ = mock_import
|
|
for mod in list(sys.modules.keys()):
|
|
if 'pretor.core.database.postgres' in mod or 'ray' in mod:
|
|
del sys.modules[mod]
|
|
|
|
from pretor.core.database.postgres import PostgresDatabase
|
|
|
|
builtins.__import__ = real_import
|
|
|
|
|
|
@patch("pretor.core.database.postgres.create_async_engine")
|
|
@patch("pretor.core.database.postgres.sessionmaker")
|
|
@patch("pretor.core.database.postgres.AuthDatabase")
|
|
@patch("pretor.core.database.postgres.ProviderDatabase")
|
|
@patch("pretor.core.database.postgres.os.environ.get")
|
|
@pytest.mark.asyncio
|
|
async def test_postgres_database(mock_env_get, mock_provider_db, mock_auth_db, mock_sessionmaker, mock_create_engine):
|
|
def env_side_effect(key):
|
|
return {
|
|
"POSTGRES_USER": "testuser",
|
|
"POSTGRES_PASSWORD": "testpassword",
|
|
"POSTGRES_HOST": "localhost",
|
|
"POSTGRES_PORT": "5432",
|
|
"POSTGRES_DB": "testdb"
|
|
}.get(key)
|
|
|
|
mock_env_get.side_effect = env_side_effect
|
|
|
|
mock_engine = MagicMock()
|
|
mock_conn = MagicMock()
|
|
from unittest.mock import AsyncMock
|
|
mock_conn.run_sync = AsyncMock()
|
|
|
|
mock_begin_ctx = MagicMock()
|
|
mock_begin_ctx.__aenter__ = AsyncMock(return_value=mock_conn)
|
|
mock_begin_ctx.__aexit__ = AsyncMock()
|
|
mock_engine.begin.return_value = mock_begin_ctx
|
|
mock_create_engine.return_value = mock_engine
|
|
|
|
db = PostgresDatabase()
|
|
|
|
mock_create_engine.assert_called_once_with(
|
|
"postgresql+asyncpg://testuser:testpassword@localhost:5432/testdb",
|
|
echo=True
|
|
)
|
|
mock_auth_db.assert_called_once()
|
|
mock_provider_db.assert_called_once()
|
|
mock_auth_db.return_value.get_user_authority = AsyncMock(return_value="test_auth")
|
|
|
|
with patch("pretor.core.database.postgres.SQLModel.metadata.create_all") as mock_create_all:
|
|
await db.init_db()
|
|
mock_conn.run_sync.assert_called_once_with(mock_create_all)
|
|
|
|
assert await db.get_user_authority(user_id="123") == "test_auth"
|