test(integration): migrate before timed owned proxy launch and gate credential probe on every worker (#46030)

* test(integration): migrate before timed owned proxy launch and gate credential probe on every worker

* test(integration): retry closed-connection samples and require both workers in the credential gate

---------

Co-authored-by: yuneng <yuneng@berri.ai>

tests/integration/_support/process.py

@@ -355,7 +355,12 @@ def _proxy_environment(
 
 
 def setup_only_proxy_run(
-    gateway: Gateway, overrides: Mapping[str, str], *, config: Path, workers: int
+    gateway: Gateway,
+    overrides: Mapping[str, str],
+    *,
+    config: Path,
+    workers: int,
+    database_setup: tuple[str, ...] = DB_PUSH,
 ) -> subprocess.CompletedProcess[str]:
     """The proxy CLI's `--skip_server_startup` pass (the image's setup step), run to completion with the
     environment an owned proxy gets."""
@@ -368,7 +373,7 @@ def setup_only_proxy_run(
             str(config),
             "--num_workers",
             str(workers),
-            *DB_PUSH,
+            *database_setup,
             "--skip_server_startup",
         ),
         cwd=_proxy_root(),

tests/integration/management/test_credential_federation_values.py

@@ -10,7 +10,7 @@ from functools import partial
 from itertools import product
 from pathlib import Path
 from types import MappingProxyType
-from typing import Final, TypeVar
+from typing import Final, Protocol, TypeVar, cast
 from urllib.parse import parse_qs
 
 import anthropic
@@ -50,6 +50,11 @@ from tests.integration._support.wire import Reply, Request, Wire, wire_server
 
 T = TypeVar("T")
 
+
+class _ClientNetworkStream(Protocol):
+    def get_extra_info(self, info: str) -> object | None: ...
+
+
 SOURCES: Final = ("token_file", "secret_reference", "internal_issuer", "keycloak", "environment")
 CLIENTS: Final = ("chat", "messages", "messages_stream")
 
@@ -771,6 +776,62 @@ def _workers(owned: OwnedProxy) -> tuple[psutil.Process, ...]:
     return tuple(child for child in psutil.Process(owned.process.pid).children() if _is_worker(child))
 
 
+def _worker_pids(owned: OwnedProxy) -> frozenset[int]:
+    return frozenset(child.pid for child in _workers(owned))
+
+
+def _client_address(response: httpx.Response) -> object | None:
+    stream: Final = cast(_ClientNetworkStream | None, response.extensions.get("network_stream"))
+    return stream.get_extra_info("client_addr") if stream is not None else None
+
+
+def _credential_read_on_worker(owned: OwnedProxy, name: str, _: int) -> int | None:
+    try:
+        with httpx.Client(
+            base_url=owned.gateway.client.base_url,
+            timeout=15,
+            trust_env=False,
+            limits=httpx.Limits(max_connections=1),
+        ) as client:
+            before: Final = client.get("/health/integration-readiness")
+            credential: Final = client.get(
+                f"/credentials/by_name/{name}",
+                headers={"Authorization": f"Bearer {owned.gateway.key}"},
+            )
+            after: Final = client.get("/health/integration-readiness")
+            addresses: Final = tuple(_client_address(response) for response in (before, credential, after))
+    except (httpx.TransportError, OSError):
+        return None
+    before_pid: Final = cast(str | None, before.headers.get("x-integration-worker-pid")) or ""
+    after_pid: Final = cast(str | None, after.headers.get("x-integration-worker-pid")) or ""
+    if (
+        before.status_code != 200
+        or credential.status_code != 200
+        or after.status_code != 200
+        or not before_pid.isdecimal()
+        or not after_pid.isdecimal()
+        or before_pid != after_pid
+        or (
+            any(address is not None for address in addresses)
+            and not all(address == addresses[0] for address in addresses)
+        )
+    ):
+        return None
+    return int(before_pid)
+
+
+def _every_worker_holds(owned: OwnedProxy, name: str) -> bool:
+    with ThreadPoolExecutor(max_workers=8) as pool:
+        samples: Final = tuple(pool.map(partial(_credential_read_on_worker, owned, name), range(16)))
+    holders: Final = frozenset(pid for pid in samples if pid is not None)
+    workers: Final = _worker_pids(owned)
+    return len(workers) == 2 and workers <= holders
+
+
+def _loaded_on_every_worker(owned: OwnedProxy, name: str) -> None:
+    assert eventually(lambda: _every_worker_holds(owned, name), lambda held: held, seconds=70)
+
+
 @pytest.mark.timeout(240)
 def test_worker_kill_mid_credential_burst(gateway: Gateway, tmp_path: Path) -> None:
     directory: Final = tmp_path.resolve()
@@ -1043,6 +1104,7 @@ def test_connection_probe_names_the_missing_id(federation: FederationRig) -> Non
     with federation.owned.gateway.scenario() as scenario:
         misconfigured: Final = _misconfigured(federation, scenario, _MISCONFIGURED["inline_token_without_org"])
         _deployments_visible(federation.owned.gateway, (misconfigured.model,))
+        _loaded_on_every_worker(federation.owned, misconfigured.credential)
         response: Final = federation.owned.gateway.request(
             "POST",
             "/health/test_connection",

tests/integration/management/test_fallback_management_public_names.py

@@ -11,9 +11,22 @@ import psutil
 import pytest
 from pydantic import JsonValue
 
-from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value
+from tests.integration._support.client import (
+    Gateway,
+    Scenario,
+    eventually,
+    gateway_from_environment,
+    object_value,
+    string_value,
+)
 from tests.integration._support.database import read_rows
-from tests.integration._support.process import OwnedProxy, owned_proxy_process
+from tests.integration._support.process import (
+    MIGRATE_DEPLOY,
+    OwnedProxy,
+    owned_proxy_process,
+    setup_only_proxy_run,
+    worker_configuration,
+)
 from tests.integration._support.redis_process import owned_redis
 
 pytestmark: Final = pytest.mark.timeout(300)
@@ -21,6 +34,20 @@ pytestmark: Final = pytest.mark.timeout(300)
 PROVIDER_KEY: Final = "integration-provider-key"
 EVICTION_WARNING: Final = "config cache eviction of router_settings failed"
 REFUSED: Final = frozenset({401, 403})
+PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml"
+
+
+@pytest.fixture(scope="module")
+def migrated_database(tmp_path_factory: pytest.TempPathFactory) -> None:
+    with gateway_from_environment() as gateway:
+        completed: Final = setup_only_proxy_run(
+            gateway,
+            {},
+            config=worker_configuration(PROXY_CONFIG, tmp_path_factory.mktemp("migrate")),
+            workers=1,
+            database_setup=MIGRATE_DEPLOY,
+        )
+    assert completed.returncode == 0, completed.stdout[-3000:] + completed.stderr[-3000:]
 
 
 @pytest.fixture
@@ -334,6 +361,7 @@ def test_context_window_and_content_policy_types_persist(gateway: Gateway) -> No
         assert policy in covered, covered
 
 
+@pytest.mark.usefixtures("migrated_database")
 def test_fallback_write_survives_a_redis_outage(gateway: Gateway, tmp_path: Path) -> None:
     with owned_redis(tmp_path) as cache:
         overrides: Final = {"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)}
@@ -357,6 +385,7 @@ def test_fallback_write_survives_a_redis_outage(gateway: Gateway, tmp_path: Path
                 cache.start()
 
 
+@pytest.mark.usefixtures("migrated_database")
 def test_fallback_write_survives_a_worker_kill(gateway: Gateway, tmp_path: Path) -> None:
     with owned_proxy_process(gateway, tmp_path, {}, workers=2, database_setup=()) as owned:
         with owned.gateway.scenario() as scenario: