push

Author: hassiebp

tests/test_core_sdk.py

@@ -1783,6 +1783,8 @@ def test_create_trace_sampling_zero():
 def test_mask_function():
     def mask_func(data):
         if isinstance(data, dict):
+            if "should_raise" in data:
+                raise
             return {k: "MASKED" for k in data}
         elif isinstance(data, str):
             return "MASKED"
@@ -1832,19 +1834,13 @@ def mask_func(data):
     assert fetched_span["input"] == {"data": "MASKED"}
     assert fetched_span["output"] == "MASKED"
 
-    # Test with faulty mask function
-    def faulty_mask_func(data):
-        raise Exception("Masking error")
-
-    langfuse = Langfuse(mask=faulty_mask_func)
-
     # Create a root span with trace properties
     with langfuse.start_as_current_span(name="test-span") as root_span:
-        root_span.update_trace(name="test_trace", input={"sensitive": "data"})
+        root_span.update_trace(name="test_trace", input={"should_raise": "data"})
         # Get trace ID for later use
         trace_id = root_span.trace_id
         # Add output to the trace
-        root_span.update_trace(output={"more": "sensitive"})
+        root_span.update_trace(output={"should_raise": "sensitive"})
 
     # Ensure data is sent
     langfuse.flush()