trustgraph/tests/unit/test_gateway/test_dispatch_mux.py
cybermaggedon 89be656990
Release/v1.2 (#457)
* Bump setup.py versions for 1.1

* PoC MCP server (#419)

* Very initial MCP server PoC for TrustGraph

* Put service on port 8000

* Add MCP container and packages to buildout

* Update docs for API/CLI changes in 1.0 (#421)

* Update some API basics for the 0.23/1.0 API change

* Add MCP container push (#425)

* Add command args to the MCP server (#426)

* Host and port parameters

* Added websocket arg

* More docs

* MCP client support (#427)

- MCP client service
- Tool request/response schema
- API gateway support for mcp-tool
- Message translation for tool request & response
- Make mcp-tool using configuration service for information
  about where the MCP services are.

* Feature/react call mcp (#428)

Key Features

  - MCP Tool Integration: Added core MCP tool support with ToolClientSpec and ToolClient classes
  - API Enhancement: New mcp_tool method for flow-specific tool invocation
  - CLI Tooling: New tg-invoke-mcp-tool command for testing MCP integration
  - React Agent Enhancement: Fixed and improved multi-tool invocation capabilities
  - Tool Management: Enhanced CLI for tool configuration and management

Changes

  - Added MCP tool invocation to API with flow-specific integration
  - Implemented ToolClientSpec and ToolClient for tool call handling
  - Updated agent-manager-react to invoke MCP tools with configurable types
  - Enhanced CLI with new commands and improved help text
  - Added comprehensive documentation for new CLI commands
  - Improved tool configuration management

Testing

  - Added tg-invoke-mcp-tool CLI command for isolated MCP integration testing
  - Enhanced agent capability to invoke multiple tools simultaneously

* Test suite executed from CI pipeline (#433)

* Test strategy & test cases

* Unit tests

* Integration tests

* Extending test coverage (#434)

* Contract tests

* Testing embeedings

* Agent unit tests

* Knowledge pipeline tests

* Turn on contract tests

* Increase storage test coverage (#435)

* Fixing storage and adding tests

* PR pipeline only runs quick tests

* Empty configuration is returned as empty list, previously was not in response (#436)

* Update config util to take files as well as command-line text (#437)

* Updated CLI invocation and config model for tools and mcp (#438)

* Updated CLI invocation and config model for tools and mcp

* CLI anomalies

* Tweaked the MCP tool implementation for new model

* Update agent implementation to match the new model

* Fix agent tools, now all tested

* Fixed integration tests

* Fix MCP delete tool params

* Update Python deps to 1.2

* Update to enable knowledge extraction using the agent framework (#439)

* Implement KG extraction agent (kg-extract-agent)

* Using ReAct framework (agent-manager-react)
 
* ReAct manager had an issue when emitting JSON, which conflicts which ReAct manager's own JSON messages, so refactored ReAct manager to use traditional ReAct messages, non-JSON structure.
 
* Minor refactor to take the prompt template client out of prompt-template so it can be more readily used by other modules. kg-extract-agent uses this framework.

* Migrate from setup.py to pyproject.toml (#440)

* Converted setup.py to pyproject.toml

* Modern package infrastructure as recommended by py docs

* Install missing build deps (#441)

* Install missing build deps (#442)

* Implement logging strategy (#444)

* Logging strategy and convert all prints() to logging invocations

* Fix/startup failure (#445)

* Fix loggin startup problems

* Fix logging startup problems (#446)

* Fix logging startup problems (#447)

* Fixed Mistral OCR to use current API (#448)

* Fixed Mistral OCR to use current API

* Added PDF decoder tests

* Fix Mistral OCR ident to be standard pdf-decoder (#450)

* Fix Mistral OCR ident to be standard pdf-decoder

* Correct test

* Schema structure refactor (#451)

* Write schema refactor spec

* Implemented schema refactor spec

* Structure data mvp (#452)

* Structured data tech spec

* Architecture principles

* New schemas

* Updated schemas and specs

* Object extractor

* Add .coveragerc

* New tests

* Cassandra object storage

* Trying to object extraction working, issues exist

* Validate librarian collection (#453)

* Fix token chunker, broken API invocation (#454)

* Fix token chunker, broken API invocation (#455)

* Knowledge load utility CLI (#456)

* Knowledge loader

* More tests
2025-08-18 20:56:09 +01:00

171 lines
No EOL
5.1 KiB
Python

"""
Tests for Gateway Dispatch Mux
"""
import pytest
import asyncio
from unittest.mock import MagicMock, AsyncMock
from trustgraph.gateway.dispatch.mux import Mux, MAX_QUEUE_SIZE
class TestMux:
"""Test cases for Mux class"""
def test_mux_initialization(self):
"""Test Mux initialization"""
mock_dispatcher_manager = MagicMock()
mock_ws = MagicMock()
mock_running = MagicMock()
mux = Mux(
dispatcher_manager=mock_dispatcher_manager,
ws=mock_ws,
running=mock_running
)
assert mux.dispatcher_manager == mock_dispatcher_manager
assert mux.ws == mock_ws
assert mux.running == mock_running
assert isinstance(mux.q, asyncio.Queue)
assert mux.q.maxsize == MAX_QUEUE_SIZE
@pytest.mark.asyncio
async def test_mux_destroy_with_websocket(self):
"""Test Mux destroy method with websocket"""
mock_dispatcher_manager = MagicMock()
mock_ws = AsyncMock()
mock_running = MagicMock()
mux = Mux(
dispatcher_manager=mock_dispatcher_manager,
ws=mock_ws,
running=mock_running
)
# Call destroy
await mux.destroy()
# Verify running.stop was called
mock_running.stop.assert_called_once()
# Verify websocket close was called
mock_ws.close.assert_called_once()
@pytest.mark.asyncio
async def test_mux_destroy_without_websocket(self):
"""Test Mux destroy method without websocket"""
mock_dispatcher_manager = MagicMock()
mock_running = MagicMock()
mux = Mux(
dispatcher_manager=mock_dispatcher_manager,
ws=None,
running=mock_running
)
# Call destroy
await mux.destroy()
# Verify running.stop was called
mock_running.stop.assert_called_once()
# No websocket to close
@pytest.mark.asyncio
async def test_mux_receive_valid_message(self):
"""Test Mux receive method with valid message"""
mock_dispatcher_manager = MagicMock()
mock_ws = AsyncMock()
mock_running = MagicMock()
mux = Mux(
dispatcher_manager=mock_dispatcher_manager,
ws=mock_ws,
running=mock_running
)
# Mock message with valid JSON
mock_msg = MagicMock()
mock_msg.json.return_value = {
"request": {"type": "test"},
"id": "test-id-123",
"service": "test-service"
}
# Call receive
await mux.receive(mock_msg)
# Verify json was called
mock_msg.json.assert_called_once()
@pytest.mark.asyncio
async def test_mux_receive_message_without_request(self):
"""Test Mux receive method with message missing request field"""
mock_dispatcher_manager = MagicMock()
mock_ws = AsyncMock()
mock_running = MagicMock()
mux = Mux(
dispatcher_manager=mock_dispatcher_manager,
ws=mock_ws,
running=mock_running
)
# Mock message without request field
mock_msg = MagicMock()
mock_msg.json.return_value = {
"id": "test-id-123"
}
# receive method should handle the RuntimeError internally
# Based on the code, it seems to catch exceptions
await mux.receive(mock_msg)
mock_ws.send_json.assert_called_once_with({"error": "Bad message"})
@pytest.mark.asyncio
async def test_mux_receive_message_without_id(self):
"""Test Mux receive method with message missing id field"""
mock_dispatcher_manager = MagicMock()
mock_ws = AsyncMock()
mock_running = MagicMock()
mux = Mux(
dispatcher_manager=mock_dispatcher_manager,
ws=mock_ws,
running=mock_running
)
# Mock message without id field
mock_msg = MagicMock()
mock_msg.json.return_value = {
"request": {"type": "test"}
}
# receive method should handle the RuntimeError internally
await mux.receive(mock_msg)
mock_ws.send_json.assert_called_once_with({"error": "Bad message"})
@pytest.mark.asyncio
async def test_mux_receive_invalid_json(self):
"""Test Mux receive method with invalid JSON"""
mock_dispatcher_manager = MagicMock()
mock_ws = AsyncMock()
mock_running = MagicMock()
mux = Mux(
dispatcher_manager=mock_dispatcher_manager,
ws=mock_ws,
running=mock_running
)
# Mock message with invalid JSON
mock_msg = MagicMock()
mock_msg.json.side_effect = ValueError("Invalid JSON")
# receive method should handle the ValueError internally
await mux.receive(mock_msg)
mock_msg.json.assert_called_once()
mock_ws.send_json.assert_called_once_with({"error": "Invalid JSON"})