#!/usr/bin/env python3
"""
Test client for the Jupyter Notebook MCP Server
Demonstrates how to use all the available tools.
"""
import asyncio
from fastmcp import Client
async def test_notebook_operations():
"""Test all notebook operations with the MCP server."""
# Connect to the MCP server
async with Client("main.py") as client:
print("🚀 Connected to Jupyter Notebook MCP Server")
# List available tools
tools = await client.list_tools()
print(f"\n📋 Available tools: {[tool.name for tool in tools]}")
notebook_path = "test_notebook.ipynb"
# 1. Get notebook info
print(f"\n📖 Getting info for {notebook_path}...")
info = await client.call_tool("get_notebook_info", {"notebook_path": notebook_path})
print(f"Notebook info: {info[0].text}")
# 2. Read all cells
print(f"\n📚 Reading all cells from {notebook_path}...")
cells = await client.call_tool("read_notebook_cells", {"notebook_path": notebook_path})
print(f"Found {len(cells)} cells")
# 3. Read only code cells
print(f"\n💻 Reading only code cells...")
code_cells = await client.call_tool("read_notebook_cells", {
"notebook_path": notebook_path,
"cell_type": "code"
})
print(f"Found {len(code_cells)} code cells")
# 4. Add a new cell
print(f"\n➕ Adding a new cell...")
new_cell_content = """# New cell added by MCP server
print("Hello from the MCP server!")
import datetime
print(f"Current time: {datetime.datetime.now()}")"""
add_result = await client.call_tool("add_cell_to_notebook", {
"notebook_path": notebook_path,
"cell_content": new_cell_content,
"cell_type": "code"
})
print(f"Add result: {add_result[0].text}")
# 5. Execute a specific cell
print(f"\n⚡ Executing cell 1 (first code cell)...")
exec_result = await client.call_tool("execute_notebook_cell", {
"notebook_path": notebook_path,
"cell_index": 1,
"timeout": 10
})
print(f"Execution result: {exec_result[0].text}")
# 6. Execute the entire notebook
print(f"\n🔄 Executing entire notebook...")
full_exec_result = await client.call_tool("execute_entire_notebook", {
"notebook_path": notebook_path,
"timeout_per_cell": 10,
"stop_on_error": False
})
print(f"Full execution result: {full_exec_result[0].text}")
# 7. Final info check
print(f"\n📊 Final notebook info...")
final_info = await client.call_tool("get_notebook_info", {"notebook_path": notebook_path})
print(f"Final info: {final_info[0].text}")
if __name__ == "__main__":
print("🧪 Testing Jupyter Notebook MCP Server")
asyncio.run(test_notebook_operations())
