Skip to content

Commit d0ab4de

Browse files
authored
fix: return 404 in artifact creation on non existing model version (kubeflow#1720)
* fix: return 404 in artifact creation on non existing model version Signed-off-by: Alessio Pragliola <[email protected]> * fix: expecting no id in response to 404 Signed-off-by: Alessio Pragliola <[email protected]> * fix: support to python 3.9 to fuzz tests Signed-off-by: Alessio Pragliola <[email protected]> --------- Signed-off-by: Alessio Pragliola <[email protected]>
1 parent bf4a769 commit d0ab4de

File tree

2 files changed

+24
-7
lines changed

2 files changed

+24
-7
lines changed

clients/python/tests/fuzz_api/model_registry/test_mr_stateless.py

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import logging
22
import secrets
33
import time
4-
from typing import Any, Callable
4+
from typing import Any, Callable, Optional
55

66
import pytest
77
import requests
@@ -87,7 +87,7 @@ def build_artifact_payload(artifact_type: str, uri_prefix: str, state: str, name
8787
return payload
8888

8989

90-
def validate_artifact_response(response: requests.Response, expected_payload: dict[str, Any]) -> str:
90+
def validate_artifact_response(response: requests.Response, expected_payload: dict[str, Any]) -> Optional[str]:
9191
"""Validate artifact creation response and return the artifact ID.
9292
9393
Args:
@@ -101,9 +101,18 @@ def validate_artifact_response(response: requests.Response, expected_payload: di
101101
AssertionError: If validation fails
102102
"""
103103
# Check response status
104-
assert response.status_code in {200, 201}, f"Expected 200 or 201, got {response.status_code}: {response.text}"
104+
assert response.status_code in {200, 201, 404}, f"Expected 200, 201, or 404, got {response.status_code}: {response.text}"
105105

106106
response_json = response.json()
107+
108+
# Handle error responses (404)
109+
if response.status_code == 404:
110+
assert "message" in response_json, "Error response should contain 'message' field"
111+
assert "not found" in response_json["message"].lower(), f"Error message should contain 'not found', got: {response_json['message']}"
112+
# For 404 responses, we don't return an ID since the artifact wasn't created
113+
return None
114+
115+
# Handle success responses (200, 201)
107116
assert response_json.get("id"), "Response body should contain 'id'"
108117

109118
# Validate response matches payload
@@ -284,8 +293,9 @@ def test_post_model_version_artifacts(auth_headers: dict, artifact_type: str, ur
284293
# Validate response and get artifact ID
285294
artifact_id = validate_artifact_response(response, payload)
286295

287-
# Cleanup after successful creation
288-
cleanup_artifacts(artifact_id)
296+
# Cleanup after successful creation (only if artifact was created)
297+
if artifact_id is not None:
298+
cleanup_artifacts(artifact_id)
289299

290300

291301
@pytest.mark.fuzz
@@ -316,8 +326,9 @@ def test_post_experiment_run_artifacts(auth_headers: dict, artifact_type: str, u
316326
# Validate response and get artifact ID
317327
artifact_id = validate_artifact_response(response, payload)
318328

319-
# Cleanup artifacts
320-
cleanup_artifacts(artifact_id)
329+
# Cleanup artifacts (only if artifact was created)
330+
if artifact_id is not None:
331+
cleanup_artifacts(artifact_id)
321332

322333
# Cleanup experiment and run
323334
cleanup_experiment_and_run(auth_headers=auth_headers, experiment_id=experiment_id, experiment_run_id=experiment_run_id, verify_tls=verify_ssl)

internal/core/artifact.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,12 @@ func (b *ModelRegistryService) upsertArtifact(artifact *openapi.Artifact, parent
314314
}
315315

316316
func (b *ModelRegistryService) UpsertModelVersionArtifact(artifact *openapi.Artifact, parentResourceId string) (*openapi.Artifact, error) {
317+
// Validate that the ModelVersion exists before creating the artifact
318+
_, err := b.GetModelVersionById(parentResourceId)
319+
if err != nil {
320+
return nil, err
321+
}
322+
317323
return b.upsertArtifact(artifact, &parentResourceId)
318324
}
319325

0 commit comments

Comments
 (0)