wehub-resource-sync
103 lines
3.5 KiB
Python
103 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from yuxi.repositories import knowledge_file_repository as repository_module
|
|
from yuxi.repositories.knowledge_file_repository import KnowledgeFileRepository
|
|
from yuxi.storage.postgres.models_knowledge import KnowledgeBase, KnowledgeFile
|
|
|
|
pytestmark = [pytest.mark.asyncio, pytest.mark.unit]
|
|
|
|
|
|
class _AsyncSessionContext:
|
|
def __init__(self, db):
|
|
self.db = db
|
|
|
|
async def __aenter__(self):
|
|
return self.db
|
|
|
|
async def __aexit__(self, exc_type, *_args):
|
|
if exc_type is None:
|
|
await self.db.commit()
|
|
else:
|
|
await self.db.rollback()
|
|
return False
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def knowledge_session(monkeypatch):
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(KnowledgeBase.__table__.create)
|
|
await conn.run_sync(KnowledgeFile.__table__.create)
|
|
|
|
session_factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
async with session_factory() as session:
|
|
monkeypatch.setattr(
|
|
repository_module.pg_manager,
|
|
"get_async_session_context",
|
|
lambda: _AsyncSessionContext(session),
|
|
)
|
|
yield session
|
|
|
|
await engine.dispose()
|
|
|
|
|
|
async def test_exists_by_filename_matches_active_file_exactly(knowledge_session):
|
|
existing_name = (
|
|
"google_drive/shared_drives/engineering/serving-runtime/dsid_e4ff04ebc2a14c1982abc4987753790c__playbook.txt"
|
|
)
|
|
knowledge_session.add_all(
|
|
[
|
|
KnowledgeBase(kb_id="kb_1", name="KB 1", description="", kb_type="milvus"),
|
|
KnowledgeBase(kb_id="kb_2", name="KB 2", description="", kb_type="milvus"),
|
|
KnowledgeFile(
|
|
file_id="file_active",
|
|
kb_id="kb_1",
|
|
filename=existing_name,
|
|
status="uploaded",
|
|
is_folder=False,
|
|
),
|
|
KnowledgeFile(
|
|
file_id="file_other_kb",
|
|
kb_id="kb_2",
|
|
filename=existing_name,
|
|
status="uploaded",
|
|
is_folder=False,
|
|
),
|
|
KnowledgeFile(
|
|
file_id="file_failed",
|
|
kb_id="kb_1",
|
|
filename="failed.txt",
|
|
status="failed",
|
|
is_folder=False,
|
|
),
|
|
KnowledgeFile(
|
|
file_id="folder_same_name",
|
|
kb_id="kb_1",
|
|
filename="folder",
|
|
status="done",
|
|
is_folder=True,
|
|
),
|
|
KnowledgeFile(
|
|
file_id="legacy_file",
|
|
kb_id="kb_1",
|
|
filename="legacy.txt",
|
|
status="indexed",
|
|
is_folder=None,
|
|
),
|
|
]
|
|
)
|
|
await knowledge_session.commit()
|
|
|
|
repo = KnowledgeFileRepository()
|
|
|
|
assert await repo.exists_by_filename(kb_id="kb_1", filename=existing_name) is True
|
|
assert await repo.exists_by_filename(kb_id="kb_1", filename=existing_name.upper()) is False
|
|
assert await repo.exists_by_filename(kb_id="kb_1", filename="failed.txt") is False
|
|
assert await repo.exists_by_filename(kb_id="kb_1", filename="folder") is False
|
|
assert await repo.exists_by_filename(kb_id="kb_1", filename="legacy.txt") is True
|
|
assert await repo.exists_by_filename(kb_id="missing", filename=existing_name) is False
|