simplify tests

Author: shellmayr

tests/integrations/langchain/test_langchain.py

@@ -1724,9 +1724,6 @@ def test_langchain_response_model_extraction(
     llm_output_model,
     expected_model,
 ):
-    from langchain_core.outputs import LLMResult
-    from langchain_core.messages import AIMessageChunk
-
     sentry_init(
         integrations=[LangchainIntegration(include_prompts=True)],
         traces_sample_rate=1.0,
@@ -1792,79 +1789,3 @@ def test_langchain_response_model_extraction(
         assert llm_span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == expected_model
     else:
         assert SPANDATA.GEN_AI_RESPONSE_MODEL not in llm_span.get("data", {})
-
-
-@pytest.mark.parametrize(
-    "missing_attribute",
-    [
-        "message",
-        "response_metadata",
-        "generation_info",
-        "llm_output",
-    ],
-)
-def test_langchain_response_model_extraction_missing_attributes(
-    sentry_init,
-    capture_events,
-    missing_attribute,
-):
-    from langchain_core.messages import AIMessageChunk
-
-    sentry_init(
-        integrations=[LangchainIntegration(include_prompts=True)],
-        traces_sample_rate=1.0,
-        send_default_pii=True,
-    )
-    events = capture_events()
-
-    callback = SentryLangchainCallback(max_span_map_size=100, include_prompts=True)
-
-    run_id = "test-missing-attr-uuid"
-    serialized = {"_type": "openai-chat", "model_name": "gpt-3.5-turbo"}
-    prompts = ["Test prompt"]
-
-    with start_transaction():
-        callback.on_llm_start(
-            serialized=serialized,
-            prompts=prompts,
-            run_id=run_id,
-            invocation_params={"model": "gpt-3.5-turbo"},
-        )
-
-        generation = Mock()
-        generation.text = "Test response"
-
-        if missing_attribute != "message":
-            message_mock = Mock()
-            message_mock.response_metadata.get.return_value = None
-            if missing_attribute == "response_metadata":
-                delattr(message_mock, "response_metadata")
-            generation.message = message_mock
-
-        if missing_attribute != "generation_info":
-            generation_info_mock = Mock()
-            generation_info_mock.get.return_value = None
-            generation.generation_info = generation_info_mock
-
-        response = Mock()
-        response.generations = [[generation]]
-
-        if missing_attribute != "llm_output":
-            llm_output_mock = Mock()
-            llm_output_mock.get.return_value = None
-            response.llm_output = llm_output_mock
-
-        callback.on_llm_end(response=response, run_id=run_id)
-
-    assert len(events) > 0
-    tx = events[0]
-    assert tx["type"] == "transaction"
-
-    llm_spans = [
-        span for span in tx.get("spans", []) if span.get("op") == "gen_ai.pipeline"
-    ]
-    assert len(llm_spans) > 0
-
-    llm_span = llm_spans[0]
-
-    assert SPANDATA.GEN_AI_RESPONSE_MODEL not in llm_span.get("data", {})