Add an E2E test for Simple Calculator Galileo observability example

examples/observability/simple_calculator_observability/tests/test_simple_calc_observability.py

@@ -27,6 +27,8 @@
 from nat.test.utils import run_workflow
 
 if typing.TYPE_CHECKING:
+    import galileo.log_streams
+    import galileo.projects
     import langsmith.client
 
 
@@ -136,17 +138,46 @@ async def test_langsmith_full_workflow(config_dir: Path,
 
     await run_workflow(config=config, question=question, expected_answer=expected_answer)
 
-    done = False
     runlist = []
     deadline = time.time() + 10
-    while not done and time.time() < deadline:
+    while len(runlist) == 0 and time.time() < deadline:
         # Wait for traces to be ingested
         await asyncio.sleep(0.5)
         runs = langsmith_client.list_runs(project_name=langsmith_project_name, is_root=True)
         runlist = [run for run in runs]
-        if len(runlist) > 0:
-            done = True
 
-    assert done, "Timed out waiting for LangSmith run to be ingested"
     # Since we have a newly created project, the above workflow should have created exactly one root run
     assert len(runlist) == 1
+
+
+@pytest.mark.integration
+@pytest.mark.usefixtures("galileo_api_key")
+async def test_galileo_full_workflow(config_dir: Path,
+                                     galileo_project: "galileo.projects.Project",
+                                     galileo_log_stream: "galileo.log_streams.LogStream",
+                                     question: str,
+                                     expected_answer: str):
+    config_file = config_dir / "config-galileo.yml"
+    config = load_config(config_file)
+    config.general.telemetry.tracing["galileo"].project = galileo_project.name
+    config.general.telemetry.tracing["galileo"].logstream = galileo_log_stream.name
+
+    await run_workflow(config=config, question=question, expected_answer=expected_answer)
+
+    import galileo.search
+
+    sessions = []
+    deadline = time.time() + 10
+    while len(sessions) == 0 and time.time() < deadline:
+        # Wait for traces to be ingested
+        await asyncio.sleep(0.5)
+        results = galileo.search.get_sessions(project_id=galileo_project.id, log_stream_id=galileo_log_stream.id)
+        sessions = results.records or []
+
+    assert len(sessions) == 1
+
+    traces = galileo.search.get_traces(project_id=galileo_project.id, log_stream_id=galileo_log_stream.id)
+    assert len(traces.records) == 1
+
+    spans = galileo.search.get_spans(project_id=galileo_project.id, log_stream_id=galileo_log_stream.id)
+    assert len(spans.records) > 1

packages/nvidia_nat_test/src/nat/test/plugin.py

@@ -27,8 +27,9 @@
 import pytest_asyncio
 
 if typing.TYPE_CHECKING:
+    import galileo.log_streams
+    import galileo.projects
     import langsmith.client
-
     from docker.client import DockerClient
 
 
@@ -290,16 +291,53 @@ def langsmith_client_fixture(langsmith_api_key: str, fail_missing: bool) -> "lan
         pytest.skip(reason=reason)
 
 
+@pytest.fixture(name="project_name")
+def project_name_fixture() -> str:
+    # Create a unique project name for each test run
+    return f"nat-e2e-test-{time.time()}-{random.random()}"
+
+
 @pytest.fixture(name="langsmith_project_name")
-def langsmith_project_name_fixture(langsmith_client: "langsmith.client.Client") -> Generator[str]:
-    # Createa a unique project name for each test run
-    project_name = f"nat-e2e-test-{time.time()}-{random.random()}"
+def langsmith_project_name_fixture(langsmith_client: "langsmith.client.Client", project_name: str) -> Generator[str]:
     langsmith_client.create_project(project_name)
     yield project_name
 
     langsmith_client.delete_project(project_name=project_name)
 
 
+@pytest.fixture(name="galileo_api_key", scope='session')
+def galileo_api_key_fixture(fail_missing: bool):
+    """
+    Use for integration tests that require a Galileo API key.
+    """
+    yield require_env_variables(
+        varnames=["GALILEO_API_KEY"],
+        reason="Galileo integration tests require the `GALILEO_API_KEY` environment variable to be defined.",
+        fail_missing=fail_missing)
+
+
+@pytest.fixture(name="galileo_project")
+def galileo_project_fixture(galileo_api_key: str, fail_missing: bool,
+                            project_name: str) -> Generator["galileo.projects.Project"]:
+    try:
+        import galileo.projects
+        project = galileo.projects.create_project(name=project_name)
+        yield project
+
+        galileo.projects.delete_project(id=project.id)
+    except ImportError:
+        reason = "Galileo integration tests require the `galileo` package to be installed."
+        if fail_missing:
+            raise RuntimeError(reason)
+        pytest.skip(reason=reason)
+
+
+@pytest.fixture(name="galileo_log_stream")
+def galileo_log_stream_fixture(galileo_project: "galileo.projects.Project") -> "galileo.log_streams.LogStream":
+    import galileo.log_streams
+    return galileo.log_streams.create_log_stream(project_id=galileo_project.id, name="test")
+
+
 @pytest.fixture(name="require_docker", scope='session')
 def require_docker_fixture(fail_missing: bool) -> "DockerClient":
     """

pyproject.toml

@@ -223,8 +223,10 @@ exclude = ["packages/compat"]
 # Dependency groups are only for developers to aid in managing dependencies local to a dev machine.
 dev = [
   "asgi-lifespan~=2.1",
+  "galileo~=1.27",
   "httpx-sse~=0.4",
   "ipython~=8.31",
+  "langsmith", # this is currently a transitive dependency of nat, not specifying a version to avoid conflicts
   "myst-parser~=4.0",
   "nbconvert", # Version determined by jupyter
   "nbsphinx~=0.9",