ci(e2e): Fix navigation e2e test - make it more flexible (#237)

Author: ioanarm

tests/conftest.py

@@ -1,5 +1,7 @@
 import pytest
 import os
+import asyncio
+import gc
 from dotenv import load_dotenv
 from mcp.client.sse import sse_client
 from mcp.client.stdio import stdio_client
@@ -20,6 +22,16 @@ def anyio_backend():
     return "asyncio"
 
 
+@pytest.fixture(autouse=True)
+async def cleanup_sessions():
+    """Clean up any lingering HTTP sessions after each test."""
+    yield
+    # Force garbage collection to clean up any unclosed sessions
+    gc.collect()
+    # Give a brief moment for cleanup
+    await asyncio.sleep(0.01)
+
+
 @pytest.fixture
 def mcp_transport():
     return os.environ.get("MCP_TRANSPORT", DEFAULT_MCP_TRANSPORT)

tests/loki_test.py

@@ -12,7 +12,7 @@
 from conftest import models
 from utils import (
     get_converted_tools,
-    llm_tool_call_sequence,
+    flexible_tool_call,
 )
 
 pytestmark = pytest.mark.anyio
@@ -30,7 +30,7 @@ async def test_loki_logs_tool(model: str, mcp_client: ClientSession):
     ]
 
     # 1. List datasources
-    messages = await llm_tool_call_sequence(
+    messages = await flexible_tool_call(
         model, messages, tools, mcp_client, "list_datasources"
     )
     datasources_response = messages[-1].content
@@ -39,8 +39,9 @@ async def test_loki_logs_tool(model: str, mcp_client: ClientSession):
     print(f"\nFound Loki datasource: {loki_ds['name']} (uid: {loki_ds['uid']})")
 
     # 2. Query logs
-    messages = await llm_tool_call_sequence(
-        model, messages, tools, mcp_client, "query_loki_logs", {"datasourceUid": loki_ds["uid"], "logql": "{container=\"mcp-grafana-grafana-1\"}"}
+    messages = await flexible_tool_call(
+        model, messages, tools, mcp_client, "query_loki_logs",
+        required_params={"datasourceUid": loki_ds["uid"]}
     )
 
     # 3. Final LLM response
@@ -66,7 +67,7 @@ async def test_loki_container_labels(model: str, mcp_client: ClientSession):
     ]
 
     # 1. List datasources
-    messages = await llm_tool_call_sequence(
+    messages = await flexible_tool_call(
         model, messages, tools, mcp_client, "list_datasources"
     )
     datasources_response = messages[-1].content
@@ -75,9 +76,9 @@ async def test_loki_container_labels(model: str, mcp_client: ClientSession):
     print(f"\nFound Loki datasource: {loki_ds['name']} (uid: {loki_ds['uid']})")
 
     # 2. List label values for 'container'
-    messages = await llm_tool_call_sequence(
+    messages = await flexible_tool_call(
         model, messages, tools, mcp_client, "list_loki_label_values",
-        {"datasourceUid": loki_ds["uid"], "labelName": "container"}
+        required_params={"datasourceUid": loki_ds["uid"], "labelName": "container"}
     )
 
     # 3. Final LLM response

tests/navigation_test.py

@@ -1,3 +1,4 @@
+import json
 import pytest
 from langevals import expect
 from langevals_langevals.llm_boolean import (
@@ -6,11 +7,13 @@
 )
 from litellm import Message, acompletion
 from mcp import ClientSession
+from mcp.types import TextContent
 
 from conftest import models
 from utils import (
     get_converted_tools,
     llm_tool_call_sequence,
+    flexible_tool_call,
 )
 
 pytestmark = pytest.mark.anyio
@@ -149,31 +152,27 @@ async def test_generate_deeplink_with_time_range(model: str, mcp_client: ClientS
 
 @pytest.mark.parametrize("model", models)
 @pytest.mark.flaky(max_runs=3)
-async def test_generate_deeplink_with_custom_params(model: str, mcp_client: ClientSession):
+async def test_generate_deeplink_with_query_params(model: str, mcp_client: ClientSession):
     tools = await get_converted_tools(mcp_client)
-    prompt = "Generate a dashboard deeplink for 'test-uid' with custom variables"
+    prompt = "Use the generate_deeplink tool to create a dashboard link for UID 'test-uid' with var-datasource=prometheus and refresh=30s as query parameters"
 
     messages = [
         Message(role="system", content="You are a helpful assistant."),
         Message(role="user", content=prompt),
     ]
 
-    messages = await llm_tool_call_sequence(
+    # Use flexible tool call with required parameters
+    messages = await flexible_tool_call(
         model, messages, tools, mcp_client, "generate_deeplink",
-        {
-            "resourceType": "dashboard",
-            "dashboardUid": "test-uid",
-            "queryParams": {
-                "var-datasource": "prometheus",
-                "refresh": "30s"
-            }
-        }
+        required_params={"resourceType": "dashboard", "dashboardUid": "test-uid"}
     )
 
     response = await acompletion(model=model, messages=messages, tools=tools)
     content = response.choices[0].message.content
 
-    assert "var-datasource=prometheus" in content, f"Expected custom parameters, got: {content}"
+    # Verify both specific query parameters are in the final URL
+    assert "var-datasource=prometheus" in content, f"Expected var-datasource=prometheus in URL, got: {content}"
+    assert "refresh=30s" in content, f"Expected refresh=30s in URL, got: {content}"
 
     custom_params_checker = CustomLLMBooleanEvaluator(
         settings=CustomLLMBooleanSettings(

tests/utils.py

@@ -68,4 +68,50 @@ async def llm_tool_call_sequence(
 
 async def get_converted_tools(mcp_client):
     tools = await mcp_client.list_tools()
-    return [convert_tool(t) for t in tools.tools]
+    return [convert_tool(t) for t in tools.tools]
+
+
+async def flexible_tool_call(model, messages, tools, mcp_client, expected_tool_name, required_params=None):
+    """
+    Make a flexible tool call that only checks essential parameters.
+    Returns updated messages list.
+
+    Args:
+        model: The LLM model to use
+        messages: Current conversation messages
+        tools: Available tools
+        mcp_client: MCP client session
+        expected_tool_name: Name of the tool we expect to be called
+        required_params: Dict of essential parameters to check (optional)
+
+    Returns:
+        Updated messages list including tool call and result
+    """
+    response = await acompletion(model=model, messages=messages, tools=tools)
+
+    # Check that a tool call was made
+    assert response.choices[0].message.tool_calls is not None, f"Expected tool call for {expected_tool_name}"
+    assert len(response.choices[0].message.tool_calls) >= 1, f"Expected at least one tool call for {expected_tool_name}"
+
+    tool_call = response.choices[0].message.tool_calls[0]
+    assert tool_call.function.name == expected_tool_name, f"Expected {expected_tool_name} tool, got {tool_call.function.name}"
+
+    arguments = json.loads(tool_call.function.arguments)
+
+    # Check required parameters if specified
+    if required_params:
+        for key, expected_value in required_params.items():
+            assert key in arguments, f"Expected parameter '{key}' in tool arguments"
+            if expected_value is not None:
+                assert arguments[key] == expected_value, f"Expected {key}='{expected_value}', got {key}='{arguments.get(key)}'"
+
+    # Call the tool to verify it works
+    result = await mcp_client.call_tool(tool_call.function.name, arguments)
+    assert len(result.content) == 1
+    assert isinstance(result.content[0], TextContent)
+
+    # Add both the tool call and result to message history
+    messages.append(response.choices[0].message)
+    messages.append(Message(role="tool", tool_call_id=tool_call.id, content=result.content[0].text))
+    
+    return messages