DocsGPT/tests/conftest.py

from unittest.mock import Mock

import pytest
from application.core.settings import settings


@pytest.fixture
def mock_llm():
    llm = Mock()
    llm.gen_stream = Mock()
    llm._supports_tools = True
    llm._supports_structured_output = Mock(return_value=False)
    llm.__class__.__name__ = "MockLLM"
    return llm


@pytest.fixture
def mock_llm_handler():
    handler = Mock()
    handler.process_message_flow = Mock()
    return handler


@pytest.fixture
def mock_retriever():
    retriever = Mock()
    retriever.search = Mock(
        return_value=[
            {"text": "Test document 1", "filename": "doc1.txt", "source": "test"},
            {"text": "Test document 2", "title": "doc2.txt", "source": "test"},
        ]
    )
    return retriever


@pytest.fixture
def mock_mongo_db(monkeypatch):
    fake_collection = FakeMongoCollection()
    fake_db = {
        "agents": fake_collection,
        "user_tools": fake_collection,
        "memories": fake_collection,
    }
    fake_client = {settings.MONGO_DB_NAME: fake_db}

    monkeypatch.setattr(
        "application.core.mongo_db.MongoDB.get_client", lambda: fake_client
    )
    return fake_client


@pytest.fixture
def sample_chat_history():
    return [
        {"prompt": "What is Python?", "response": "Python is a programming language."},
        {"prompt": "Tell me more.", "response": "Python is known for simplicity."},
    ]


@pytest.fixture
def sample_tool_call():
    return {
        "tool_name": "test_tool",
        "call_id": "123",
        "action_name": "test_action",
        "arguments": {"arg1": "value1"},
        "result": "Tool executed successfully",
    }


@pytest.fixture
def decoded_token():
    return {"sub": "test_user", "email": "test@example.com"}


@pytest.fixture
def log_context():
    from application.logging import LogContext

    context = LogContext(
        endpoint="test_endpoint",
        activity_id="test_activity",
        user="test_user",
        api_key="test_key",
        query="test query",
    )
    return context


class FakeMongoCollection:
    def __init__(self):
        self.docs = {}

    def find_one(self, query, projection=None):
        if "key" in query:
            return self.docs.get(query["key"])
        if "_id" in query:
            return self.docs.get(str(query["_id"]))
        if "user" in query:
            for doc in self.docs.values():
                if doc.get("user") == query["user"]:
                    return doc
        return None

    def find(self, query, projection=None):
        results = []
        if "_id" in query and "$in" in query["_id"]:
            for doc_id in query["_id"]["$in"]:
                doc = self.docs.get(str(doc_id))
                if doc:
                    results.append(doc)
        elif "user" in query:
            for doc in self.docs.values():
                if doc.get("user") == query["user"]:
                    if "status" in query:
                        if doc.get("status") == query["status"]:
                            results.append(doc)
                    else:
                        results.append(doc)
        return results

    def insert_one(self, doc):
        doc_id = doc.get("_id", len(self.docs))
        self.docs[str(doc_id)] = doc
        return Mock(inserted_id=doc_id)

    def update_one(self, query, update, upsert=False):
        return Mock(modified_count=1)

    def delete_one(self, query):
        return Mock(deleted_count=1)

    def delete_many(self, query):
        return Mock(deleted_count=0)


@pytest.fixture
def mock_llm_creator(mock_llm, monkeypatch):
    monkeypatch.setattr(
        "application.llm.llm_creator.LLMCreator.create_llm", Mock(return_value=mock_llm)
    )
    return mock_llm


@pytest.fixture
def mock_llm_handler_creator(mock_llm_handler, monkeypatch):
    monkeypatch.setattr(
        "application.llm.handlers.handler_creator.LLMHandlerCreator.create_handler",
        Mock(return_value=mock_llm_handler),
    )
    return mock_llm_handler


@pytest.fixture
def agent_base_params(decoded_token):
    return {
        "endpoint": "https://api.example.com",
        "llm_name": "openai",
        "gpt_model": "gpt-4",
        "api_key": "test_api_key",
        "user_api_key": None,
        "prompt": "You are a helpful assistant.",
        "chat_history": [],
        "decoded_token": decoded_token,
        "attachments": [],
        "json_schema": None,
    }


@pytest.fixture
def mock_tool():
    tool = Mock()
    tool.execute_action = Mock(return_value="Tool result")
    tool.get_actions_metadata = Mock(
        return_value=[
            {
                "name": "test_action",
                "description": "A test action",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "param1": {"type": "string", "description": "Test parameter"}
                    },
                    "required": ["param1"],
                },
            }
        ]
    )
    return tool


@pytest.fixture
def mock_tool_manager(mock_tool, monkeypatch):
    manager = Mock()
    manager.load_tool = Mock(return_value=mock_tool)
    monkeypatch.setattr(
        "application.agents.base.ToolManager", Mock(return_value=manager)
    )
    return manager