splunk-mcp

import pytest from unittest.mock import Mock, patch, AsyncMock, MagicMock import json from datetime import datetime from splunk_mcp import ( get_splunk_connection, search_splunk, list_indexes, list_users, list_kvstore_collections, create_kvstore_collection, delete_kvstore_collection, ) # Mock data for testing MOCK_SEARCH_RESULTS = { "results": [ {"field1": "value1", "field2": "value2"}, {"field1": "value3", "field2": "value4"}, ] } MOCK_INDEXES = [ { "name": "main", "totalEventCount": "1000", "currentDBSizeMB": "100", "maxTotalDataSizeMB": "500", "earliestTime": "0", "latestTime": "1615430400", } ] MOCK_USERS = [ { "name": "admin", "realname": "Administrator", "email": "admin@example.com", "defaultApp": "search", "type": "admin", } ] MOCK_KVSTORE_COLLECTIONS = [ { "name": "test_collection", "app": "search", "fields": {"field1": "string", "field2": "number"}, "accelerated_fields": {}, "record_count": 10, } ] @pytest.fixture def mock_splunk_service(): """Fixture to create a mock Splunk service""" mock_service = Mock() # Mock search job mock_job = Mock() mock_results = Mock() mock_results.read.return_value = json.dumps(MOCK_SEARCH_RESULTS).encode('utf-8') mock_job.results.return_value = mock_results mock_service.jobs.create.return_value = mock_job # Mock indexes mock_index = MagicMock() # Use MagicMock for __getitem__ support mock_index.name = MOCK_INDEXES[0]["name"] mock_index.__getitem__.side_effect = lambda x: MOCK_INDEXES[0][x] mock_service.indexes = [mock_index] # Mock users mock_user = MagicMock() # Use MagicMock for __getitem__ support mock_user.name = MOCK_USERS[0]["name"] mock_user.__getitem__.side_effect = lambda x: MOCK_USERS[0][x] mock_user.role_entities.return_value = ["admin"] mock_service.users = [mock_user] # Mock KV store mock_collection = MagicMock() # Use MagicMock for attribute access mock_collection.name = MOCK_KVSTORE_COLLECTIONS[0]["name"] mock_collection.fields = MOCK_KVSTORE_COLLECTIONS[0]["fields"] mock_collection.accelerated_fields = MOCK_KVSTORE_COLLECTIONS[0]["accelerated_fields"] mock_collection.data = range(10) mock_app = MagicMock() # Use MagicMock for attribute access mock_app.name = "search" mock_app.kvstore = [mock_collection] mock_service.apps = [mock_app] return mock_service @pytest.mark.asyncio @patch('splunk_mcp.get_splunk_connection') async def test_search_splunk(mock_get_conn, mock_splunk_service): """Test the search_splunk function""" mock_get_conn.return_value = mock_splunk_service results = await search_splunk( search_query="search index=main", earliest_time="-1h", latest_time="now", max_results=10 ) assert results == MOCK_SEARCH_RESULTS["results"] mock_splunk_service.jobs.create.assert_called_once() search_args = mock_splunk_service.jobs.create.call_args[0][0] assert search_args == "search index=main" @pytest.mark.asyncio @patch('splunk_mcp.get_splunk_connection') async def test_list_indexes(mock_get_conn, mock_splunk_service): """Test the list_indexes function""" mock_get_conn.return_value = mock_splunk_service indexes = await list_indexes() assert len(indexes) == 1 assert indexes[0]["name"] == MOCK_INDEXES[0]["name"] assert indexes[0]["total_event_count"] == MOCK_INDEXES[0]["totalEventCount"] @pytest.mark.asyncio @patch('splunk_mcp.get_splunk_connection') async def test_list_users(mock_get_conn, mock_splunk_service): """Test the list_users function""" mock_get_conn.return_value = mock_splunk_service users = await list_users() assert len(users) == 1 assert users[0]["username"] == MOCK_USERS[0]["name"] assert users[0]["real_name"] == MOCK_USERS[0]["realname"] assert users[0]["roles"] == ["admin"] @pytest.mark.asyncio @patch('splunk_mcp.get_splunk_connection') async def test_list_kvstore_collections(mock_get_conn, mock_splunk_service): """Test the list_kvstore_collections function""" mock_get_conn.return_value = mock_splunk_service collections = await list_kvstore_collections() assert len(collections) == 1 assert collections[0]["name"] == MOCK_KVSTORE_COLLECTIONS[0]["name"] assert collections[0]["app"] == MOCK_KVSTORE_COLLECTIONS[0]["app"] assert collections[0]["fields"] == MOCK_KVSTORE_COLLECTIONS[0]["fields"] @pytest.mark.asyncio @patch('splunk_mcp.get_splunk_connection') async def test_create_kvstore_collection(mock_get_conn, mock_splunk_service): """Test the create_kvstore_collection function""" mock_get_conn.return_value = mock_splunk_service fields = {"field1": "string", "field2": "number"} result = await create_kvstore_collection( collection_name="test_collection", app_name="search", fields=fields ) assert result["name"] == MOCK_KVSTORE_COLLECTIONS[0]["name"] assert result["app"] == MOCK_KVSTORE_COLLECTIONS[0]["app"] assert result["fields"] == fields @pytest.mark.asyncio @patch('splunk_mcp.get_splunk_connection') async def test_delete_kvstore_collection(mock_get_conn, mock_splunk_service): """Test the delete_kvstore_collection function""" mock_get_conn.return_value = mock_splunk_service result = await delete_kvstore_collection( collection_name="test_collection", app_name="search" ) assert result is True