simplify test

Author: shellmayr

tests/integrations/langchain/test_langchain.py

@@ -593,53 +593,72 @@ def test_langchain_callback_list_existing_callback(sentry_init):
 
 def test_tools_integration_in_spans(sentry_init, capture_events):
     """Test that tools are properly set on spans in actual LangChain integration."""
+    global llm_type
+    llm_type = "openai-chat"
+
     sentry_init(
         integrations=[LangchainIntegration(include_prompts=False)],
         traces_sample_rate=1.0,
     )
     events = capture_events()
 
-    # Create a simple mock LLM that supports tools
-    from langchain_core.language_models.llms import LLM
-
-    class MockLLMWithTools(LLM):
-        def _call(self, prompt, stop=None, run_manager=None, **kwargs):
-            return "Mock response"
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a helpful assistant"),
+            ("user", "{input}"),
+            MessagesPlaceholder(variable_name="agent_scratchpad"),
+        ]
+    )
 
-        @property
-        def _llm_type(self):
-            return "mock_llm"
+    global stream_result_mock
+    stream_result_mock = Mock(
+        side_effect=[
+            [
+                ChatGenerationChunk(
+                    type="ChatGenerationChunk",
+                    message=AIMessageChunk(content="Simple response"),
+                ),
+            ]
+        ]
+    )
 
-    # Mock tools for testing
-    test_tools = [
-        {"name": "search", "description": "Search tool"},
-        {"name": "calculator", "description": "Math tool"},
-    ]
+    llm = MockOpenAI(
+        model_name="gpt-3.5-turbo",
+        temperature=0,
+        openai_api_key="badkey",
+    )
+    agent = create_openai_tools_agent(llm, [get_word_length], prompt)
+    agent_executor = AgentExecutor(agent=agent, tools=[get_word_length], verbose=True)
 
     with start_transaction():
-        llm = MockLLMWithTools()
-        # Simulate a call with tools
-        llm.invoke("test prompt", config={"tools": test_tools})
+        list(agent_executor.stream({"input": "Hello"}))
 
-    # Check that events were captured
+    # Check that events were captured and contain tools data
     if events:
         tx = events[0]
         spans = tx.get("spans", [])
 
         # Look for spans that should have tools data
+        tools_found = False
         for span in spans:
             span_data = span.get("data", {})
             if SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in span_data:
+                tools_found = True
                 tools_data = span_data[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS]
                 # Verify tools are in the expected format
                 assert isinstance(tools_data, (str, list))  # Could be serialized
                 if isinstance(tools_data, str):
-                    # If serialized as string, should contain tool names
-                    assert "search" in tools_data
-                    assert "calculator" in tools_data
+                    # If serialized as string, should contain tool name
+                    assert "get_word_length" in tools_data
                 else:
                     # If still a list, verify structure
-                    assert len(tools_data) == 2
-                    names = [tool.get("name") for tool in tools_data]
-                    assert "search" in names
-                    assert "calculator" in names
+                    assert len(tools_data) >= 1
+                    names = [
+                        tool.get("name")
+                        for tool in tools_data
+                        if isinstance(tool, dict)
+                    ]
+                    assert "get_word_length" in names
+
+        # Ensure we found at least one span with tools data
+        assert tools_found, "No spans found with tools data"