cBioPortal MCP Server

import pytest from unittest.mock import patch, call @pytest.fixture def mock_studies_data_page_1(): return [ { "studyId": f"study_{i}", "name": f"Study {i}", "description": f"Description {i}", } for i in range(1, 4) ] @pytest.fixture def mock_studies_data_page_2(): return [ { "studyId": f"study_{i}", "name": f"Study {i}", "description": f"Description {i}", } for i in range(4, 7) ] @pytest.fixture def mock_studies_data_last_page_less_than_pagesize(): return [ { "studyId": f"study_{i}", "name": f"Study {i}", "description": f"Description {i}", } for i in range(7, 9) ] @pytest.fixture def mock_studies_data_last_page_exact_pagesize(): return [ { "studyId": f"study_{i}", "name": f"Study {i}", "description": f"Description {i}", } for i in range(9, 12) ] @pytest.fixture def mock_cancer_types_data(): return [ {"cancerTypeId": f"type_{i}", "name": f"Cancer Type {i}"} for i in range(1, 51) ] @pytest.fixture def mock_samples_data(): return [ { "sampleId": f"sample_{i}", "patientId": f"patient_{i % 20}", "studyId": "study_1", } for i in range(1, 201) ] @pytest.fixture def mock_mutations_data_page_1(): return [ { "uniqueSampleKey": f"sample_{i}_study_1", "uniquePatientKey": f"patient_{i % 20}_study_1", "molecularProfileId": "profile_1", "sampleId": f"sample_{i}", "patientId": f"patient_{i % 20}", "studyId": "study_1", "gene": {"hugoGeneSymbol": "GENE1", "entrezGeneId": 123}, "mutationEffect": "MISSENSE", "mutationStatus": "SOMATIC", "mutationType": "SNP", "proteinChange": "p.V600E", "keyword": "V600E", } for i in range(1, 3) ] @pytest.fixture def mock_molecular_profiles_for_mutations_test(): return [ { "molecularProfileId": "brca_tcga_pan_can_atlas_2018_mutations", "studyId": "brca_tcga_pan_can_atlas_2018", "name": "BRCA TCGA PanCanAtlas 2018 Mutations", "molecularAlterationType": "MUTATION_EXTENDED", "datatype": "MAF", "showProfileInAnalysisTab": True, } ] @pytest.fixture def mock_clinical_data_page_1(): return [ { "uniqueSampleKey": f"sample_{i}_study_1", "uniquePatientKey": f"patient_{i % 10}_study_1", "clinicalAttributeId": "ATTR_1", "patientId": f"patient_{i % 10}", "sampleId": f"sample_{i}", "studyId": "study_1", "value": f"Value {i}", } for i in range(1, 3) ] @pytest.fixture def mock_clinical_attributes_data(): return [ { "clinicalAttributeId": f"ATTR_{i}", "displayName": f"Attribute {i}", "description": f"Description for Attribute {i}", "datatype": "STRING", "patientAttribute": True, } for i in range(1, 76) ] @pytest.fixture def mock_genes_data(): return [ { "entrezGeneId": i, "hugoGeneSymbol": f"GENE{i}", "type": "protein-coding", "oncogene": i % 2 == 0, "tsg": i % 2 != 0, } for i in range(1, 151) ] @pytest.fixture def mock_molecular_profiles_data_all(): return [ { "molecularProfileId": f"profile_{i}", "studyId": f"study_{(i % 5) + 1}", "name": f"Molecular Profile {i}", "molecularAlterationType": "MUTATION_EXTENDED", "datatype": "MAF", "showProfileInAnalysisTab": True, } for i in range(1, 61) ] # Define test cases for get_cancer_studies pagination cancer_studies_test_cases = [ pytest.param( "first_page", # scenario_name 0, # page_number "mock_studies_data_page_1", # mock_data_fixture_name True, # expected_has_more 3, # page_size_to_use id="first_page_full", ), pytest.param( "second_page", 1, "mock_studies_data_page_2", True, 3, id="second_page_full" ), pytest.param( "last_page_partial", 2, "mock_studies_data_last_page_less_than_pagesize", False, # API returns less than page_size, so server says no more 3, id="last_page_less_than_pagesize", ), pytest.param( "last_page_full", 3, "mock_studies_data_last_page_exact_pagesize", True, # API returns exactly page_size, so server says more might exist 3, id="last_page_exact_pagesize", ), ] @pytest.mark.asyncio @patch("cbioportal_mcp.api_client.APIClient.make_api_request") @pytest.mark.parametrize( "scenario_name, page_number, mock_data_fixture_name, expected_has_more, page_size_to_use", cancer_studies_test_cases, ) async def test_get_cancer_studies_pagination( mock_api_request, cbioportal_server_instance, request, # Pytest request fixture to dynamically get other fixtures scenario_name, page_number, mock_data_fixture_name, expected_has_more, page_size_to_use, ): server = cbioportal_server_instance # Dynamically get the mock data fixture by its name mock_data = request.getfixturevalue(mock_data_fixture_name) mock_api_request.return_value = mock_data expected_items_count = len(mock_data) result = await server.get_cancer_studies( page_number=page_number, page_size=page_size_to_use ) assert len(result["studies"]) == expected_items_count assert result["pagination"]["page"] == page_number assert result["pagination"]["has_more"] is expected_has_more mock_api_request.assert_called_with( "studies", params={ "pageNumber": page_number, "pageSize": page_size_to_use, "direction": "ASC", }, ) @pytest.mark.asyncio @patch("cbioportal_mcp.api_client.APIClient.make_api_request") async def test_get_mutations_in_gene_pagination( mock_api_request, cbioportal_server_instance, mock_mutations_data_page_1, mock_molecular_profiles_for_mutations_test, # Mock for the internal molecular profile fetch ): server = cbioportal_server_instance gene_id = "TP53" study_id = "brca_tcga_pan_can_atlas_2018" sample_list_id = f"{study_id}_all" # Required by the server method page_size = 2 # Must match mock data for has_more logic # Setup side_effect for the two internal calls: # 1. Fetch molecular profiles for the study # 2. Fetch mutations for the determined molecular profile mock_api_request.side_effect = [ mock_molecular_profiles_for_mutations_test, # First call mock_mutations_data_page_1, # Second call ] result = await server.get_mutations_in_gene( gene_id=gene_id, study_id=study_id, sample_list_id=sample_list_id, page_number=0, page_size=page_size, ) assert len(result["mutations"]) == page_size assert result["pagination"]["page"] == 0 assert result["pagination"]["has_more"] is True # Assert the first call to get molecular profiles mock_api_request.assert_any_call(f"studies/{study_id}/molecular-profiles") # Assert the second call to get mutations (assuming a mutation_profile_id was found) # This requires knowing/mocking the mutation_profile_id that would be selected. # For this test, let's assume 'brca_tcga_pan_can_atlas_2018_mutations' is selected from the mock. expected_mutation_profile_id = "brca_tcga_pan_can_atlas_2018_mutations" mock_api_request.assert_called_with( f"molecular-profiles/{expected_mutation_profile_id}/mutations", method="GET", params={ "studyId": study_id, "sampleListId": sample_list_id, "pageNumber": 0, "pageSize": page_size, "direction": "ASC", "hugoGeneSymbol": gene_id, }, ) @pytest.mark.asyncio @patch("cbioportal_mcp.api_client.APIClient.make_api_request") async def test_get_clinical_data_pagination( mock_api_request, cbioportal_server_instance, mock_clinical_data_page_1 ): server = cbioportal_server_instance study_id = "acc_tcga" page_size = 2 # Must match mock data for has_more mock_api_request.return_value = mock_clinical_data_page_1 result = await server.get_clinical_data( study_id=study_id, page_number=0, page_size=page_size ) # The server method transforms the flat list into a dict by patientId # So, len(result["clinical_data_by_patient"]) is the count of unique patients. # The mock_clinical_data_page_1 contains 2 items for 2 unique patients. assert len(result["clinical_data_by_patient"]) == 2 assert result["pagination"]["page"] == 0 assert ( result["pagination"]["has_more"] is True ) # Based on page_size vs items from API mock_api_request.assert_called_with( f"studies/{study_id}/clinical-data", method="GET", params={ "pageNumber": 0, "pageSize": page_size, "direction": "ASC", "clinicalDataType": "PATIENT", }, ) @pytest.mark.asyncio @patch("cbioportal_mcp.api_client.APIClient.make_api_request") async def test_get_molecular_profiles_pagination( mock_api_request, cbioportal_server_instance, mock_molecular_profiles_data_all, # Use a mock that represents all profiles ): server = cbioportal_server_instance study_id = "study_123" # Example study_id page_size = 2 # Server will paginate client-side from the full list # Server fetches all profiles first, then paginates them. mock_api_request.return_value = mock_molecular_profiles_data_all result = await server.get_molecular_profiles( study_id=study_id, page_number=0, page_size=page_size ) # Assuming mock_molecular_profiles_data_all has more than 'page_size' items assert len(result["molecular_profiles"]) == page_size assert result["pagination"]["page"] == 0 assert result["pagination"]["has_more"] is True # The API call should be to fetch all profiles for the study, without pagination params mock_api_request.assert_called_with(f"studies/{study_id}/molecular-profiles") @pytest.mark.asyncio @patch("cbioportal_mcp.api_client.APIClient.make_api_request") async def test_paginate_results_basic( mock_make_api_request, cbioportal_server_instance ): server = cbioportal_server_instance endpoint = "studies" page_size = 2 mock_page_1_data = [{"id": 1}, {"id": 2}] mock_page_2_data = [{"id": 3}, {"id": 4}] mock_empty_page_data = [] # Configure the mock to return different data for sequential calls mock_make_api_request.side_effect = [ mock_page_1_data, mock_page_2_data, mock_empty_page_data, # Signifies no more data ] collected_results = [] async for page_data in server.paginate_results( endpoint, params={"pageSize": page_size} ): collected_results.extend(page_data) assert collected_results == mock_page_1_data + mock_page_2_data expected_calls = [ call( endpoint, method="GET", params={"pageNumber": 0, "pageSize": page_size}, json_data=None, ), call( endpoint, method="GET", params={"pageNumber": 1, "pageSize": page_size}, json_data=None, ), call( endpoint, method="GET", params={"pageNumber": 2, "pageSize": page_size}, json_data=None, ), # This call returns empty ] mock_make_api_request.assert_has_calls(expected_calls) assert mock_make_api_request.call_count == 3 @pytest.mark.asyncio @patch("cbioportal_mcp.api_client.APIClient.make_api_request") async def test_paginate_results_empty_first_call( mock_make_api_request, cbioportal_server_instance ): server = cbioportal_server_instance endpoint = "studies" page_size = 2 mock_empty_page_data = [] mock_make_api_request.return_value = mock_empty_page_data collected_results = [] async for page_data in server.paginate_results( endpoint, params={"pageNumber": 0, "pageSize": page_size} ): collected_results.extend(page_data) assert collected_results == [] # Expect only one call, which returns empty expected_calls = [ call( endpoint, method="GET", params={"pageNumber": 0, "pageSize": page_size}, json_data=None, ), ] mock_make_api_request.assert_has_calls(expected_calls) assert mock_make_api_request.call_count == 1 @pytest.mark.asyncio @patch("cbioportal_mcp.api_client.APIClient.make_api_request") async def test_paginate_results_with_max_pages( mock_make_api_request, cbioportal_server_instance ): server = cbioportal_server_instance endpoint = "studies" page_size = 1 max_pages_to_fetch = 2 # Mock API to simulate more data than max_pages mock_page_1_data = [{"id": 1}] mock_page_2_data = [{"id": 2}] mock_page_3_data = [{"id": 3}] # This page should not be fetched mock_make_api_request.side_effect = [ mock_page_1_data, mock_page_2_data, mock_page_3_data, ] collected_results = [] async for page_data in server.paginate_results( endpoint, params={"pageNumber": 0, "pageSize": page_size}, max_pages=max_pages_to_fetch, ): collected_results.extend(page_data) assert collected_results == mock_page_1_data + mock_page_2_data expected_calls = [ call( endpoint, method="GET", params={"pageNumber": 0, "pageSize": page_size}, json_data=None, ), call( endpoint, method="GET", params={"pageNumber": 1, "pageSize": page_size}, json_data=None, ), # No call for pageNumber 2 because max_pages is 2 ] mock_make_api_request.assert_has_calls(expected_calls) assert mock_make_api_request.call_count == 2 # Should only call API twice @pytest.mark.asyncio @patch("cbioportal_mcp.api_client.APIClient.make_api_request") async def test_paginate_results_last_page_partial( mock_make_api_request, cbioportal_server_instance ): server = cbioportal_server_instance endpoint = "studies" page_size = 3 mock_page_1_data = [{"id": 1}, {"id": 2}, {"id": 3}] # Full page mock_page_2_partial_data = [ {"id": 4}, {"id": 5}, ] # Partial page, less than page_size mock_make_api_request.side_effect = [ mock_page_1_data, mock_page_2_partial_data, ] collected_results = [] async for page_data in server.paginate_results( endpoint, params={"pageNumber": 0, "pageSize": page_size} ): collected_results.extend(page_data) assert collected_results == mock_page_1_data + mock_page_2_partial_data expected_calls = [ call( endpoint, method="GET", params={"pageNumber": 0, "pageSize": page_size}, json_data=None, ), call( endpoint, method="GET", params={"pageNumber": 1, "pageSize": page_size}, json_data=None, ), ] mock_make_api_request.assert_has_calls(expected_calls) assert mock_make_api_request.call_count == 2 # API called for page 0 and page 1