from unittest import mock from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import FastAPI from fastapi.testclient import TestClient from prometheus.app.api.routes import repository from prometheus.app.entity.repository import Repository from prometheus.app.exception_handler import register_exception_handlers app = FastAPI() register_exception_handlers(app) app.include_router(repository.router, prefix="/repository", tags=["repository"]) client = TestClient(app) @pytest.fixture def mock_service(): service = mock.MagicMock() app.state.service = service yield service def test_upload_repository(mock_service): mock_service["repository_service"].clone_github_repo = AsyncMock(return_value="/mock/path") mock_service["repository_service"].get_repository_by_url_and_commit_id = AsyncMock( return_value=None ) mock_service["repository_service"].create_new_repository = AsyncMock(return_value=1) mock_service["knowledge_graph_service"].build_and_save_knowledge_graph = AsyncMock( return_value=0 ) response = client.post( "/repository/upload", json={ "github_token": "mock_token", "https_url": "https://github.com/Pantheon-temple/Prometheus", }, ) assert response.status_code == 200 assert response.json() == { "code": 200, "message": "success", "data": {"repository_id": 1}, } def test_upload_repository_at_commit(mock_service): mock_service["repository_service"].clone_github_repo = AsyncMock(return_value="/mock/path") mock_service["repository_service"].get_repository_by_url_and_commit_id = AsyncMock( return_value=None ) mock_service["repository_service"].create_new_repository = AsyncMock(return_value=1) mock_service["knowledge_graph_service"].build_and_save_knowledge_graph = AsyncMock( return_value=0 ) response = client.post( "/repository/upload/", json={ "github_token": "mock_token", "https_url": "https://github.com/Pantheon-temple/Prometheus", "commit_id": "0c554293648a8705769fa53ec896ae24da75f4fc", }, ) assert response.status_code == 200 def test_create_branch_and_push(mock_service): # Mock git_repo git_repo_mock = MagicMock() git_repo_mock.create_and_push_branch = AsyncMock(return_value=None) # Let repository_service.get_repository return the mocked git_repo mock_service["repository_service"].get_repository.return_value = git_repo_mock mock_service["repository_service"].get_repository_by_id = AsyncMock( return_value=Repository( id=1, url="https://github.com/fake/repo.git", commit_id=None, playground_path="/path/to/playground", kg_root_node_id=0, user_id=None, kg_max_ast_depth=100, kg_chunk_size=1000, kg_chunk_overlap=100, ) ) response = client.post( "/repository/create-branch-and-push/", json={ "repository_id": 1, "branch_name": "new_branch", "commit_message": "Initial commit on new branch", "patch": "mock_patch_content", }, ) assert response.status_code == 200 @mock.patch("prometheus.app.api.routes.repository.delete_repository_memory") def test_delete(mock_delete_memory, mock_service): # Mock the delete_repository_memory to return success mock_delete_memory.return_value = {"code": 200, "message": "success", "data": None} mock_service["repository_service"].get_repository_by_id = AsyncMock( return_value=Repository( id=1, url="https://github.com/fake/repo.git", commit_id=None, playground_path="/path/to/playground", kg_root_node_id=0, user_id=None, kg_max_ast_depth=100, kg_chunk_size=1000, kg_chunk_overlap=100, ) ) mock_service["knowledge_graph_service"].clear_kg = AsyncMock(return_value=None) mock_service["repository_service"].clean_repository.return_value = None mock_service["repository_service"].delete_repository = AsyncMock(return_value=None) response = client.delete( "repository/delete", params={ "repository_id": 1, }, ) assert response.status_code == 200 def test_list(mock_service): mock_service["repository_service"].get_all_repositories = AsyncMock( return_value=[ Repository( id=1, url="https://github.com/fake/repo.git", commit_id=None, playground_path="/path/to/playground", kg_root_node_id=0, user_id=None, kg_max_ast_depth=100, kg_chunk_size=1000, kg_chunk_overlap=100, ) ] ) response = client.get("repository/list/") assert response.status_code == 200 assert response.json() == { "code": 200, "message": "success", "data": [ { "id": 1, "url": "https://github.com/fake/repo.git", "commit_id": None, "is_working": False, "user_id": None, "kg_max_ast_depth": 100, "kg_chunk_size": 1000, "kg_chunk_overlap": 100, } ], }