86 lines
2.6 KiB
Python
86 lines
2.6 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 "kilostar.core.postgres_database.postgres" in mod or "ray" in mod:
|
|
del sys.modules[mod]
|
|
|
|
from kilostar.core.postgres_database.postgres import PostgresDatabase # noqa: E402
|
|
|
|
builtins.__import__ = real_import
|
|
|
|
|
|
@patch("kilostar.core.postgres_database.postgres.create_async_engine")
|
|
@patch("kilostar.core.postgres_database.postgres.sessionmaker")
|
|
@patch("kilostar.core.postgres_database.postgres.AuthDatabase")
|
|
@patch("kilostar.core.postgres_database.postgres.ProviderDatabase")
|
|
@patch("kilostar.core.postgres_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(
|
|
"kilostar.core.postgres_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"
|
|
# noqa: E402
|