diff --git a/tests/unit/drivers/observability/test_griptape_cloud_observability_driver.py b/tests/unit/drivers/observability/test_griptape_cloud_observability_driver.py index 98b5ea043e..f559a90779 100644 --- a/tests/unit/drivers/observability/test_griptape_cloud_observability_driver.py +++ b/tests/unit/drivers/observability/test_griptape_cloud_observability_driver.py @@ -1,3 +1,6 @@ +import os +from uuid import UUID + import pytest from opentelemetry.sdk.trace import Event, ReadableSpan from opentelemetry.trace import SpanContext, Status, StatusCode @@ -10,9 +13,23 @@ class TestGriptapeCloudObservabilityDriver: @pytest.fixture() def driver(self): - return GriptapeCloudObservabilityDriver( - base_url="http://base-url:1234", api_key="api-key", structure_run_id="structure-run-id" - ) + environ = { + "GT_CLOUD_BASE_URL": "http://base-url:1234", + "GT_CLOUD_API_KEY": "api-key", + "GT_CLOUD_STRUCTURE_RUN_ID": "structure-run-id", + } + original_environ = {} + for key in environ: + original_environ[key] = environ.get(key) + os.environ[key] = environ[key] + + yield GriptapeCloudObservabilityDriver() + + for key, value in original_environ.items(): + if value is None: + del os.environ[key] + else: + os.environ[key] = value @pytest.fixture(autouse=True) def mock_span_exporter_class(self, mocker): @@ -39,6 +56,10 @@ def test_init(self, mock_span_exporter_class, mock_span_exporter): mock_span_exporter.export.assert_not_called() + def test_init_raises_when_structure_run_is_none(self): + with pytest.raises(ValueError, match="structure_run_id must be set"): + GriptapeCloudObservabilityDriver(structure_run_id=None) + def test_context_manager_pass(self, driver, mock_span_exporter): expected_spans = ExpectedSpans(spans=[ExpectedSpan(name="main", parent=None, status_code=StatusCode.OK)]) @@ -145,6 +166,19 @@ def method(self, word: str): mock_span_exporter.export.assert_called_with(expected_spans) mock_span_exporter.export.reset_mock() + def test_get_span_id(self, driver): + assert driver.get_span_id() is None + with driver: + # Span ID's returned from GriptapeCloudObservabilityDriver should be valid UUIDs + assert self._is_valid_uuid(driver.get_span_id()) + + def _is_valid_uuid(self, val: str) -> bool: + try: + UUID(str(val)) + return True + except ValueError: + return False + class TestGriptapeCloudObservabilityDriverSpanExporter: @pytest.fixture() diff --git a/tests/unit/drivers/observability/test_open_telemetry_observability_driver.py b/tests/unit/drivers/observability/test_open_telemetry_observability_driver.py index 44da647270..b903fb1c94 100644 --- a/tests/unit/drivers/observability/test_open_telemetry_observability_driver.py +++ b/tests/unit/drivers/observability/test_open_telemetry_observability_driver.py @@ -188,3 +188,10 @@ def func(word: str): assert mock_span_exporter.export.call_count == 1 mock_span_exporter.export.assert_called_with(expected_spans) mock_span_exporter.export.reset_mock() + + def test_get_span_id(self, driver): + assert driver.get_span_id() is None + with driver: + span_id = driver.get_span_id() + assert span_id is not None + assert isinstance(span_id, str)