From 959211cdde2014df156de38c58159f1a94206289 Mon Sep 17 00:00:00 2001 From: Agochar Bhatia <37423596+agocharbhatia@users.noreply.github.com> Date: Wed, 4 Nov 2026 19:04:58 -0500 Subject: [PATCH] Reject failed reference uploads before registration --- src/openlayer/lib/data/reference_dataset.py | 4 +- tests/test_data_uploads.py | 173 ++++++++++++++++++++ 2 files changed, 176 insertions(+), 1 deletion(-) create mode 100644 tests/test_data_uploads.py diff --git a/src/openlayer/lib/data/reference_dataset.py b/src/openlayer/lib/data/reference_dataset.py index 45b3d76c..5ed60e15 100644 --- a/src/openlayer/lib/data/reference_dataset.py +++ b/src/openlayer/lib/data/reference_dataset.py @@ -48,11 +48,13 @@ def upload_reference_dataframe( tar.add(tmp_dir, arcname=os.path.basename("reference_dataset")) # Upload to storage - uploader.upload( + response = uploader.upload( file_path=tar_file_path, object_name=object_name, presigned_url_response=presigned_url_response, ) + if response.status_code >= 300 or response.status_code < 200: + raise ValueError(f"Failed to upload file to storage: {response.text}") # Notify the backend client.inference_pipelines.update( diff --git a/tests/test_data_uploads.py b/tests/test_data_uploads.py new file mode 100644 index 00000000..d6ae571d --- /dev/null +++ b/tests/test_data_uploads.py @@ -0,0 +1,173 @@ +"""Keep dataset registration behind a successful storage upload.""" + +from __future__ import annotations + +import json +import socket +from typing import Any, Protocol +from dataclasses import field, dataclass + +import httpx +import pandas as pd # pyright: ignore[reportMissingTypeStubs] +import pytest +import requests + +from openlayer import Openlayer +from openlayer.lib import data +from openlayer.types.inference_pipelines import data_stream_params + + +# ----------------------------- Helper functions ----------------------------- # +@dataclass +class Boundary: + status: int + events: list[str] = field(default_factory=lambda: list[str]()) + registrations: list[dict[str, Any]] = field(default_factory=lambda: list[dict[str, Any]]()) + transport_error: bool = False + + def api(self, request: httpx.Request) -> httpx.Response: + """Keep generated request construction and response parsing in the path.""" + assert request.url.host == "api.openlayer.invalid" + if request.method == "POST" and request.url.path == "/storage/presigned-url": + self.events.append("presign") + object_name = request.url.params["objectName"] + return httpx.Response( + 200, + json={ + "url": "https://storage.openlayer.invalid/upload", + "storageUri": "s3://synthetic-bucket/" + object_name, + "fields": {"key": object_name, "policy": "synthetic-policy"}, + }, + ) + assert (request.method, request.url.path) in { + ("PUT", "/inference-pipelines/pipeline-synthetic"), + ("POST", "/inference-pipelines/pipeline-synthetic/data"), + } + self.events.append("register") + self.registrations.append(json.loads(request.content)) + return httpx.Response(200, json={"id": "pipeline-synthetic"}) + + def storage(self, request: requests.PreparedRequest) -> requests.Response: + """Replace the external storage service after request preparation.""" + assert request.url == "https://storage.openlayer.invalid/upload" + self.events.append("storage") + if self.transport_error: + raise requests.Timeout("synthetic storage timeout") + response = requests.Response() + response.status_code = self.status + response._content = b"synthetic storage response" + response.request = request + response.url = request.url + return response + + +class BoundaryFactory(Protocol): + def __call__(self, status: int = 204, transport_error: bool = False) -> Boundary: ... + + +@pytest.fixture(autouse=True) +def forbid_network(monkeypatch: pytest.MonkeyPatch) -> None: + """Turn a missing mock into a local failure before it can open a connection.""" + + def deny(*_args: object, **_kwargs: object) -> None: + raise AssertionError("Live network access is forbidden") + + monkeypatch.setattr(socket.socket, "connect", deny) + monkeypatch.setattr(socket, "create_connection", deny) + + +@pytest.fixture +def boundary_factory(monkeypatch: pytest.MonkeyPatch) -> BoundaryFactory: + def create(status: int = 204, transport_error: bool = False) -> Boundary: + boundary = Boundary(status=status, transport_error=transport_error) + + def send(_session: requests.Session, request: requests.PreparedRequest, **_kwargs: object) -> requests.Response: + return boundary.storage(request) + + monkeypatch.setattr(requests.sessions.Session, "send", send) + return boundary + + return create + + +def invoke(boundary: Boundary, storage: data.StorageType, workflow: str) -> None: + """Use a real client so a failed upload can be observed at registration.""" + frame = pd.DataFrame({"feature": [10, 20], "prediction": [0, 1]}) + config: data_stream_params.Config = { + "class_names": ["negative", "positive"], + "feature_names": ["feature"], + "predictions_column_name": "prediction", + } + with httpx.Client(transport=httpx.MockTransport(boundary.api), trust_env=False) as http: + with Openlayer( + api_key="synthetic-not-a-credential", + base_url="https://api.openlayer.invalid", + http_client=http, + max_retries=0, + ) as client: + upload = data.upload_reference_dataframe if workflow == "reference" else data.upload_batch_inferences + upload( + client=client, + inference_pipeline_id="pipeline-synthetic", + dataset_df=frame, + config=config, + storage_type=storage, + ) + + +# --------------------------- End helper functions --------------------------- # + + +@pytest.mark.parametrize("storage", list(data.StorageType)) +@pytest.mark.parametrize("status", [200, 201, 204]) +def test_reference_success_registers_after_upload( + boundary_factory: BoundaryFactory, storage: data.StorageType, status: int +) -> None: + boundary = boundary_factory(status) + invoke(boundary, storage, "reference") + assert boundary.events == ["presign", "storage", "register"] + assert len(boundary.registrations) == 1 + assert set(boundary.registrations[0]) == {"referenceDatasetUri"} + assert boundary.registrations[0]["referenceDatasetUri"].startswith("s3://synthetic-bucket/reference_dataset_") + + +@pytest.mark.parametrize("storage", list(data.StorageType)) +@pytest.mark.parametrize("workflow", ["reference", "batch"]) +@pytest.mark.parametrize("status", [403, 503]) +def test_rejected_storage_does_not_register( + boundary_factory: BoundaryFactory, storage: data.StorageType, workflow: str, status: int +) -> None: + """Rejecting the object must prevent publication of its storage URI.""" + boundary = boundary_factory(status) + with pytest.raises(ValueError, match="Failed to upload file to storage: synthetic storage response"): + invoke(boundary, storage, workflow) + assert boundary.events == ["presign", "storage"] + assert boundary.registrations == [] + + +@pytest.mark.parametrize("storage", list(data.StorageType)) +@pytest.mark.parametrize("workflow", ["reference", "batch"]) +def test_transport_failure_does_not_register( + boundary_factory: BoundaryFactory, storage: data.StorageType, workflow: str +) -> None: + boundary = boundary_factory(transport_error=True) + with pytest.raises(requests.Timeout, match="synthetic storage timeout"): + invoke(boundary, storage, workflow) + assert boundary.events == ["presign", "storage"] + assert boundary.registrations == [] + + +@pytest.mark.parametrize("storage", list(data.StorageType)) +def test_batch_success_control(boundary_factory: BoundaryFactory, storage: data.StorageType) -> None: + boundary = boundary_factory() + invoke(boundary, storage, "batch") + assert boundary.events == ["presign", "storage", "register"] + assert len(boundary.registrations) == 1 + registration = boundary.registrations[0] + assert registration["performDataMerge"] is False + assert registration["storageUri"].endswith(".arrow") + assert registration["config"] == { + "classNames": ["negative", "positive"], + "featureNames": ["feature"], + "predictionsColumnName": "prediction", + }