test(integration): remove flaky dependencies from the core suite (ENG-3999)

tests/helpers/__init__.py

@@ -6,11 +6,13 @@
     default_labels,
     default_region,
     env,
+    run_id,
     sleep,
     unique_name,
     wait_for_sandbox_deletion,
     wait_for_sandbox_deployed,
     wait_for_volume_deletion,
+    wait_until,
 )
 
 __all__ = [
@@ -19,9 +21,11 @@
     "default_labels",
     "default_region",
     "env",
+    "run_id",
     "sleep",
     "unique_name",
     "wait_for_sandbox_deletion",
     "wait_for_sandbox_deployed",
     "wait_for_volume_deletion",
+    "wait_until",
 ]

tests/helpers/echo.py

@@ -0,0 +1,110 @@
+"""Self-hosted httpbin replacement for the proxy integration tests.
+
+The proxy suite needs an external HTTPS endpoint that echoes back the headers
+and body it received. It used to call ``httpbin.org``, which rate-limits the
+shared Blaxel egress IPs and answers ``503`` in bursts -- taking ~30 tests down
+with it on roughly half of the CI runs.
+
+Instead we host the echo ourselves: one sandbox, one node server, one public
+preview. Setup costs ~3s per pytest session and the endpoint is a real external
+HTTPS host from the sandbox's point of view, so the proxy/firewall paths under
+test are exercised exactly the same way.
+"""
+
+from __future__ import annotations
+
+from urllib.parse import urlparse
+
+from blaxel.core.sandbox import SandboxInstance
+
+from .utils import default_image, default_labels, default_region, unique_name
+
+# Serves the subset of httpbin routes the proxy suite actually used:
+# /headers, /get, /post, /put, /delete, /redirect/N, /bytes/N.
+ECHO_SERVER_SCRIPT = r"""
+const http = require("http");
+
+http.createServer((req, res) => {
+  let body = "";
+  req.on("data", (c) => (body += c));
+  req.on("end", () => {
+    if (req.url.startsWith("/redirect/")) {
+      const n = parseInt(req.url.split("/")[2] || "1", 10);
+      res.writeHead(302, { Location: n > 1 ? "/redirect/" + (n - 1) : "/get" });
+      res.end();
+      return;
+    }
+    if (req.url.startsWith("/bytes/")) {
+      const n = parseInt(req.url.split("/")[2] || "0", 10);
+      res.writeHead(200, { "Content-Type": "application/octet-stream" });
+      res.end(Buffer.alloc(n, "a"));
+      return;
+    }
+    let parsed = null;
+    try {
+      parsed = body ? JSON.parse(body) : null;
+    } catch (e) {}
+    const host = req.headers["x-forwarded-host"] || req.headers.host;
+    const proto = req.headers["x-forwarded-proto"] || "https";
+    res.writeHead(200, { "Content-Type": "application/json" });
+    res.end(
+      JSON.stringify({
+        headers: req.headers,
+        method: req.method,
+        url: proto + "://" + host + req.url,
+        data: body,
+        json: parsed,
+        args: {},
+      })
+    );
+  });
+}).listen(3000, "0.0.0.0", () => console.log("echo listening on 3000"));
+""".strip()
+
+_ECHO_PORT = 3000
+
+# Cached for the lifetime of the pytest process: every test class reuses the
+# same endpoint instead of paying the setup again.
+_echo_url: str | None = None
+
+
+async def echo_url() -> str:
+    """Return the base HTTPS URL of the shared echo server, creating it if needed.
+
+    The sandbox carries ``default_labels`` so the session-level cleanup in
+    ``tests/integration/core/conftest.py`` deletes it with everything else.
+    """
+    global _echo_url
+    if _echo_url is not None:
+        return _echo_url
+
+    sandbox = await SandboxInstance.create_if_not_exists(
+        {
+            "name": unique_name("echo"),
+            "image": default_image,
+            "region": default_region,
+            "labels": default_labels,
+        }
+    )
+    await sandbox.fs.write("/tmp/echo-server.js", ECHO_SERVER_SCRIPT)
+    await sandbox.process.exec(
+        {
+            "name": "echo-server",
+            "command": "node /tmp/echo-server.js",
+            "wait_for_ports": [_ECHO_PORT],
+        }
+    )
+    preview = await sandbox.previews.create_if_not_exists(
+        {
+            "metadata": {"name": "echo"},
+            "spec": {"port": _ECHO_PORT, "public": True},
+        }
+    )
+
+    _echo_url = preview.spec.url
+    return _echo_url
+
+
+async def echo_host() -> str:
+    """Hostname of the echo server, for firewall allow/deny lists."""
+    return urlparse(await echo_url()).hostname

tests/helpers/utils.py

@@ -13,10 +13,17 @@
 default_region = "eu-dub-1" if env == "dev" else "us-pdx-1"
 default_image = "blaxel/base-image:latest"
 
+# Unique per pytest process. CI runs of several PRs share one workspace, so the
+# end-of-session cleanup must only delete what *this* run created -- deleting by
+# ``env=integration-test`` alone tears down sandboxes a concurrent run is still
+# using, which is a large part of the suite's cross-run flakiness.
+run_id = uuid.uuid4().hex[:12]
+
 # Default labels to identify test sandboxes in the UI
 default_labels = {
     "env": "integration-test",
     "created-by": "pytest",
+    "run-id": run_id,
 }
 
 
@@ -107,6 +114,21 @@ async def wait_for_volume_deletion(volume_name: str, max_attempts: int = 30) ->
     return False
 
 
+async def wait_until(predicate, timeout: float = 10.0, interval: float = 0.1) -> bool:
+    """Poll ``predicate`` until it is true or ``timeout`` elapses.
+
+    Callbacks (watch events, log streams) usually fire in well under a second,
+    but a fixed sleep turns a slow round-trip into a test failure. Polling keeps
+    the fast path fast and the slow path green.
+    """
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if predicate():
+            return True
+        await asyncio.sleep(interval)
+    return predicate()
+
+
 def sleep(seconds: float) -> None:
     """Synchronous sleep helper."""
     time.sleep(seconds)

tests/integration/core/conftest.py

@@ -1,12 +1,42 @@
 """Pytest configuration for core integration tests."""
 
 import asyncio
+from datetime import datetime, timedelta, timezone
+
+# Orphans older than this were left behind by a crashed run, never by a live one.
+ORPHAN_MAX_AGE = timedelta(hours=2)
+
+
+def _labels(resource) -> dict:
+    metadata = getattr(resource, "metadata", None)
+    labels = getattr(metadata, "labels", None) if metadata else None
+    return getattr(labels, "additional_properties", {}) or {}
+
+
+def _is_stale_orphan(resource) -> bool:
+    """True for a pytest resource old enough that no live run still needs it."""
+    created_at = getattr(getattr(resource, "metadata", None), "created_at", None)
+    if not isinstance(created_at, str):
+        return False
+    try:
+        # Timestamps come back with nanosecond precision, which fromisoformat
+        # rejects on older Pythons -- truncate to microseconds.
+        head, _, tail = created_at.partition(".")
+        created = datetime.fromisoformat(f"{head}.{tail[:6]}+00:00" if tail else f"{head}+00:00")
+    except ValueError:
+        return False
+    return datetime.now(timezone.utc) - created > ORPHAN_MAX_AGE
 
 
 def pytest_sessionfinish(session, exitstatus):
-    """Clean up all test sandboxes after the test session ends.
+    """Clean up the sandboxes this run created.
 
     With pytest-xdist, this only runs on the master node after all workers finish.
+
+    Only sandboxes tagged with this run's ``run-id`` are deleted, plus stale
+    orphans from runs that crashed. CI runs several pull requests against the
+    same workspace: deleting every ``env=integration-test`` sandbox would tear
+    down sandboxes a concurrent run is still using.
     """
     # Skip cleanup on worker nodes (pytest-xdist)
     # Workers have workerinput attribute, master doesn't
@@ -15,46 +45,34 @@ def pytest_sessionfinish(session, exitstatus):
 
     from blaxel.core.client.client import client
     from blaxel.core.sandbox import SandboxInstance
-    from blaxel.core.volume import VolumeInstance
+    from tests.helpers import run_id
+
+    def is_ours(resource) -> bool:
+        labels = _labels(resource)
+        if labels.get("run-id") == run_id:
+            return True
+        return labels.get("created-by") == "pytest" and _is_stale_orphan(resource)
 
     async def cleanup_test_resources():
-        """Delete all sandboxes and volumes with test labels."""
+        """Delete this run's sandboxes, plus stale orphans."""
         # Reset client for cleanup
         client._async_client = None
 
         print("\n🧹 Cleaning up test resources...")
 
-        # Clean up sandboxes with test labels
+        # Volumes are not swept here: the list endpoint returns LiteVolumeMetadata,
+        # which carries no labels, so there is no way to tell ours apart without a
+        # GET per volume. Volume tests delete what they create in their own
+        # class-level fixtures.
         try:
             sandboxes = await SandboxInstance.list()
-            for sb in sandboxes:
-                labels = sb.metadata.labels
-                # Labels are stored in additional_properties of MetadataLabels object
-                if labels is not None:
-                    props = getattr(labels, "additional_properties", {}) or {}
-                    if props.get("env") == "integration-test":
-                        try:
-                            await sb.delete()
-                        except Exception:
-                            pass
         except Exception as e:
             print(f"  Error listing sandboxes: {e}")
-
-        # Clean up volumes with test labels
-        try:
-            volumes = await VolumeInstance.list()
-            for vol in volumes:
-                labels = vol.metadata.labels if hasattr(vol, "metadata") and vol.metadata else None
-                # Labels are stored in additional_properties of MetadataLabels object
-                if labels is not None:
-                    props = getattr(labels, "additional_properties", {}) or {}
-                    if props.get("env") == "integration-test":
-                        try:
-                            await vol.delete()  # type: ignore[attr-defined]
-                        except Exception:
-                            pass
-        except Exception as e:
-            print(f"  Error listing volumes: {e}")
+            sandboxes = []
+        await asyncio.gather(
+            *(sb.delete() for sb in sandboxes if is_ours(sb)),
+            return_exceptions=True,
+        )
 
         # Close the client
         if client._async_client is not None:

tests/integration/core/sandbox/proxy/helpers.py

@@ -10,6 +10,7 @@
 import os
 
 from blaxel.core.client.types import Unset
+from tests.helpers.echo import echo_url
 
 # The proxy/network routing feature is only available in specific regions, so
 # these tests override ``tests.helpers.default_region`` (which points at
@@ -20,8 +21,9 @@
 PROXY_HELPER_SCRIPT = r"""
 const https = require("https");
 const tls = require("tls");
+const RESPONSE_HEADERS_MARKER = "__RESPONSE_HEADERS__";
 const method = process.argv[2] || "GET";
-const targetUrl = process.argv[3] || "https://httpbin.org/headers";
+const targetUrl = process.argv[3];
 const extraHeaders = process.argv[4] ? JSON.parse(process.argv[4]) : {};
 const bodyData = process.argv[5] || null;
 const proxyUrl = process.env.HTTPS_PROXY || process.env.https_proxy ||
@@ -41,7 +43,11 @@
   }
   const req = https.request(opts, (r) => {
     let d = ""; r.on("data", c => d += c);
-    r.on("end", () => { process.stdout.write(d); process.exit(0); });
+    r.on("end", () => {
+      process.stdout.write(d);
+      process.stdout.write("\n" + RESPONSE_HEADERS_MARKER + JSON.stringify(r.headers));
+      process.exit(0);
+    });
   });
   req.on("error", (e) => { process.stderr.write("REQ ERR: " + e.message + "\n"); process.exit(1); });
   if (bodyData) req.write(bodyData);
@@ -91,14 +97,28 @@
 PYTHON_HELPER_SCRIPT = """
 import sys, json, requests
 method = sys.argv[1] if len(sys.argv) > 1 else "GET"
-url = sys.argv[2] if len(sys.argv) > 2 else "https://httpbin.org/headers"
+url = sys.argv[2]
 headers = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
 body = sys.argv[4] if len(sys.argv) > 4 else None
 resp = requests.request(method, url, headers=headers, data=body, timeout=30)
 print(resp.text)
+print("__RESPONSE_HEADERS__" + json.dumps(dict(resp.headers)))
 """.strip()
 
 
+# Each test sandbox gets the echo server base URL in this file, and the test
+# commands read it back with ``$(cat /tmp/echo-url)``. Commands stay plain
+# string literals this way -- no f-strings fighting with the JSON braces they
+# already contain.
+ECHO_URL_FILE = "/tmp/echo-url"
+
+
+async def write_echo_url(sandbox) -> None:
+    """Make the echo base URL readable as ``ECHO`` from commands in this sandbox."""
+    sandbox_echo_url = await echo_url()
+    await sandbox.fs.write(ECHO_URL_FILE, sandbox_echo_url)
+
+
 def not_unset(val) -> bool:
     """Return True if val is a real value (not Unset/None)."""
     return val is not None and not isinstance(val, Unset)
@@ -127,5 +147,23 @@ def parse_json_output(logs: str | None) -> dict:
     return json.loads(trimmed[start:end])
 
 
+RESPONSE_HEADERS_MARKER = "__RESPONSE_HEADERS__"
+
+
+def parse_response_headers(logs: str | None) -> dict[str, str]:
+    """Headers of the HTTP response as the client inside the sandbox saw them.
+
+    The proxy stamps ``X-Blaxel-Request-Id`` on the response it hands back, so
+    this is where that header is observable. It never reaches the upstream
+    server: the cluster gateway strips inbound ``x-blaxel-*`` headers before
+    forwarding, by design.
+    """
+    marker_at = (logs or "").rfind(RESPONSE_HEADERS_MARKER)
+    if marker_at == -1:
+        raise ValueError(f"No response headers in output: {(logs or '')[:200]}")
+    raw = logs[marker_at + len(RESPONSE_HEADERS_MARKER) :].strip()
+    return lowercase_keys(json.loads(raw))
+
+
 def lowercase_keys(obj: dict[str, str]) -> dict[str, str]:
     return {k.lower(): v for k, v in obj.items()}