From c86300491fab61af69224d42064d587bdadd02a4 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 3 Nov 2026 14:27:15 -0800 Subject: [PATCH] test(integration): use rc 108 database read helper --- .../observability/test_presidio_masking_endpoints.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/integration/observability/test_presidio_masking_endpoints.py b/tests/integration/observability/test_presidio_masking_endpoints.py index 33a436be8d5..a2e4f875856 100644 --- a/tests/integration/observability/test_presidio_masking_endpoints.py +++ b/tests/integration/observability/test_presidio_masking_endpoints.py @@ -11,6 +11,7 @@ from typing import Final import pytest import yaml from integration._support.client import Gateway, eventually, gateway_from_environment, list_value, object_value +from integration._support.database import read_rows from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter @@ -479,7 +480,7 @@ def _ask(rig: Rig, model: str, surface: Surface, guardrail: str, prompt: str) -> def _spend_row(rig: Rig, model: str) -> dict[str, JsonValue]: rows: Final = eventually( - lambda: rig.gateway.rows('SELECT response, metadata FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)), + lambda: read_rows('SELECT response, metadata FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)), lambda found: len(found) == 1, seconds=70, )