enh(tracing): Add `output_dir` arg and unit tests. (#11)

src/surf_spot_finder/tracing.py

@@ -23,28 +23,32 @@ class JsonFileSpanExporter(SpanExporter):
         pass
 
 
-def get_tracer_provider(project_name: str, json_tracer: bool) -> TracerProvider:
+def get_tracer_provider(
+    project_name: str, json_tracer: bool, output_dir: str = "telemetry_output"
+) -> TracerProvider:
     """
     Create a tracer_provider based on the selected mode.
 
     Args:
         project_name: Name of the project for tracing
         json_tracer: Whether to use the custom JSON file exporter (True) or Phoenix (False)
+        output_dir: The directory where the telemetry output will be stored.
+            Only used if `json_tracer=True`.
+            Defaults to "telemetry_output".
 
     Returns:
         TracerProvider: The configured tracer provider
     """
     if json_tracer:
-        local_folder: str = "telemetry_output"
-        if not os.path.exists(local_folder):
-            os.makedirs(local_folder)
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
         timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
 
         tracer_provider = TracerProvider()
         trace.set_tracer_provider(tracer_provider)
 
         json_file_exporter = JsonFileSpanExporter(
-            file_name=f"{local_folder}/{project_name}-{timestamp}.json"
+            file_name=f"{output_dir}/{project_name}-{timestamp}.json"
         )
         span_processor = SimpleSpanProcessor(json_file_exporter)
         tracer_provider.add_span_processor(span_processor)

tests/unit/test_unit_tracing.py

@@ -0,0 +1,35 @@
+from unittest.mock import patch, MagicMock
+
+import pytest
+
+from surf_spot_finder.tracing import get_tracer_provider, setup_tracing
+
+
+@pytest.mark.parametrize("json_tracer", [True, False])
+def test_get_tracer_provider(tmp_path, json_tracer):
+    mock_trace = MagicMock()
+    mock_tracer_provider = MagicMock()
+    mock_register = MagicMock()
+
+    with (
+        patch("surf_spot_finder.tracing.trace", mock_trace),
+        patch("surf_spot_finder.tracing.TracerProvider", mock_tracer_provider),
+        patch("surf_spot_finder.tracing.register", mock_register),
+    ):
+        get_tracer_provider(
+            project_name="test_project",
+            json_tracer=json_tracer,
+            output_dir=tmp_path / "telemetry",
+        )
+        assert (tmp_path / "telemetry").exists() == json_tracer
+        if json_tracer:
+            mock_trace.set_tracer_provider.assert_called_once_with(
+                mock_tracer_provider.return_value
+            )
+        else:
+            mock_register.assert_called_once_with(project_name="test_project")
+
+
+def test_invalid_agent_type():
+    with pytest.raises(ValueError, match="agent_type must be one of"):
+        setup_tracing(MagicMock(), "invalid_agent_type")