remora/crates/remora-contracts/tests/schema_contract.rs

491 lines
15 KiB
Rust

use std::{collections::BTreeSet, fs};
use remora_contracts::{
ActivationResponseDto, AttemptActivateRequestDto, AttemptCheckpointRequestDto,
AttemptOutcomeRequestDto, AttemptProgressRequestDto, AttemptRenewRequestDto,
AttemptResolvedSourceRequestDto, AttemptStagingMultipartRequestDto,
AttemptStagingObjectRequestDto, AttemptStagingWriteIntentRequestDto, AttemptStopRequestDto,
ClaimRequestDto, ClaimResponseDto, JobPhaseProjectionDto, LeaseProjectionDto,
};
use serde::{Serialize, de::DeserializeOwned};
use serde_json::Value;
#[test]
fn api_contract_documents_parse_and_local_references_resolve() {
let openapi_source = fs::read_to_string("../../docs/api/openapi.yaml").unwrap();
let openapi: Value = serde_yaml_ng::from_str(&openapi_source).unwrap();
assert_eq!(
openapi.pointer("/openapi").and_then(Value::as_str),
Some("3.2.0")
);
assert_local_references_resolve(&openapi);
let events_source = fs::read_to_string("../../docs/api/events.schema.json").unwrap();
let events: Value = serde_json::from_str(&events_source).unwrap();
assert_eq!(
events.pointer("/$schema").and_then(Value::as_str),
Some("https://json-schema.org/draft/2020-12/schema")
);
assert_local_references_resolve(&events);
jsonschema::meta::validate(&events).unwrap();
let event_validator = jsonschema::draft202012::options()
.should_validate_formats(true)
.build(&events)
.unwrap();
let corpus: Value =
serde_json::from_str(include_str!("../../../contracts/golden/domain.json")).unwrap();
for frame in corpus["event_frames"].as_array().unwrap() {
event_validator.validate(frame).unwrap();
}
let component_schema = openapi_component_schema(&openapi);
jsonschema::meta::validate(&component_schema).unwrap();
validate_openapi_corpus(
&component_schema,
&corpus,
"JobSubmission",
"job_submissions",
);
validate_openapi_corpus(
&component_schema,
&corpus,
"AttemptOutcomeRequest",
"attempt_outcomes",
);
validate_openapi_corpus(&component_schema, &corpus, "JobFailure", "job_failures");
}
#[test]
fn worker_routes_and_dtos_conform_to_the_mutable_openapi_contract() {
let openapi: Value =
serde_yaml_ng::from_str(&fs::read_to_string("../../docs/api/openapi.yaml").unwrap())
.unwrap();
let component_schema = openapi_component_schema(&openapi);
let mutation = serde_json::json!({
"operation_id": "operation_0001",
"fence": 1,
"worker_incarnation": "worker_0001"
});
let descriptor = serde_json::json!({
"id": "descriptor_0001",
"kind": "resolved_source",
"schema_version": 1,
"byte_length": 64,
"item_count": 1,
"checksum": format!("sha256:{}", "22".repeat(32))
});
let execution = serde_json::json!({
"id": "grant_0001",
"job_id": "job_00000001",
"attempt_id": "attempt_0001",
"fence": 1,
"worker_incarnation": "worker_0001",
"source": {
"kind": "http",
"display_url": "https://example.invalid/artifact",
"result_path": "artifact.bin"
},
"sensitive_source_ref": "source_00001",
"credential_version": null,
"resolved_source_ref": null,
"target": {
"store": "primary",
"bucket": "remora-results",
"prefix": "users/alice"
},
"staging_target": {
"store": "primary",
"bucket": "remora-results",
"prefix": "users/alice/_remora/a/attempt_0001/"
},
"execution_grant_ref": "grant_0001",
"checkpoint_ref": null
});
let authority = serde_json::json!({
"attempt_id": "attempt_0001",
"fence": 1,
"worker_incarnation": "worker_0001"
});
let lease = serde_json::json!({
"expires_at": "2026-07-17T00:01:00Z",
"renew_after": "2026-07-17T00:00:30Z"
});
let claim_request = serde_json::json!({
"operation_id": "operation_0001",
"worker_incarnation": "worker_0001"
});
assert_dto_and_schema::<ClaimRequestDto>(&component_schema, "ClaimRequest", claim_request);
let claim_response = serde_json::json!({
"authority": authority,
"execution": execution
});
assert_dto_and_schema::<ClaimResponseDto>(
&component_schema,
"ClaimResponse",
claim_response.clone(),
);
assert_dto_and_schema::<AttemptRenewRequestDto>(
&component_schema,
"AttemptRenewRequest",
mutation.clone(),
);
assert_dto_and_schema::<AttemptActivateRequestDto>(
&component_schema,
"AttemptActivateRequest",
mutation.clone(),
);
assert_dto_and_schema::<AttemptStopRequestDto>(
&component_schema,
"AttemptStopRequest",
serde_json::json!({
"operation_id": "operation_0001",
"fence": 1,
"worker_incarnation": "worker_0001"
}),
);
assert_dto_and_schema::<AttemptStagingObjectRequestDto>(
&component_schema,
"AttemptStagingObjectRequest",
serde_json::json!({
"operation_id": "operation_0001",
"fence": 1,
"worker_incarnation": "worker_0001",
"object": {
"store": "primary",
"bucket": "remora-results",
"key": "users/alice/_remora/a/attempt_0001/o/object_0001"
}
}),
);
assert_dto_and_schema::<AttemptStagingWriteIntentRequestDto>(
&component_schema,
"AttemptStagingWriteIntentRequest",
serde_json::json!({
"operation_id": "operation_0001",
"fence": 1,
"worker_incarnation": "worker_0001",
"intent": {
"store": "primary",
"bucket": "remora-results",
"key": "users/alice/_remora/a/attempt_0001/o/object_0001"
}
}),
);
assert_dto_and_schema::<AttemptStagingMultipartRequestDto>(
&component_schema,
"AttemptStagingMultipartRequest",
serde_json::json!({
"operation_id": "operation_0001",
"fence": 1,
"worker_incarnation": "worker_0001",
"multipart": {
"store": "primary",
"bucket": "remora-results",
"key": "users/alice/_remora/a/attempt_0001/o/object_0001",
"upload_id": "upload_0001"
}
}),
);
assert_dto_and_schema::<ActivationResponseDto>(
&component_schema,
"ActivationResponse",
serde_json::json!({
"authority": claim_response["authority"],
"execution": claim_response["execution"],
"lease": lease.clone()
}),
);
assert_dto_and_schema::<LeaseProjectionDto>(
&component_schema,
"LeaseProjection",
lease.clone(),
);
assert_dto_and_schema::<AttemptResolvedSourceRequestDto>(
&component_schema,
"AttemptResolvedSourceRequest",
merged(
&mutation,
serde_json::json!({ "resolved_source_ref": descriptor.clone() }),
),
);
let mut checkpoint_descriptor = descriptor.clone();
checkpoint_descriptor["kind"] = Value::String("checkpoint".into());
assert_dto_and_schema::<AttemptCheckpointRequestDto>(
&component_schema,
"AttemptCheckpointRequest",
merged(
&mutation,
serde_json::json!({
"checkpoint_ref": checkpoint_descriptor,
"resolved_source_id": "descriptor_0001",
"validated_bytes": 1024,
"validated_members": 1
}),
),
);
assert_dto_and_schema::<AttemptProgressRequestDto>(
&component_schema,
"AttemptProgressRequest",
merged(
&mutation,
serde_json::json!({
"progress": {
"completed_bytes": 1024,
"total_bytes": 4096,
"completed_members": 1,
"total_members": 1
}
}),
),
);
assert_dto_and_schema::<AttemptOutcomeRequestDto>(
&component_schema,
"AttemptOutcomeRequest",
merged(
&mutation,
serde_json::json!({
"outcome": {
"kind": "failure",
"failure": {
"class": "transient_source",
"message": "temporarily unavailable"
}
}
}),
),
);
assert_dto_and_schema::<JobPhaseProjectionDto>(
&component_schema,
"JobPhaseProjection",
serde_json::json!({ "phase": "RetryWait" }),
);
let routes = [
(
"/internal/v1/jobs/{job_id}/claim",
"ClaimRequest",
"200",
Some("ClaimResponse"),
),
(
"/internal/v1/attempts/{attempt_id}/renew",
"AttemptRenewRequest",
"200",
Some("LeaseProjection"),
),
(
"/internal/v1/attempts/{attempt_id}/activate",
"AttemptActivateRequest",
"200",
Some("ActivationResponse"),
),
(
"/internal/v1/attempts/{attempt_id}/stop",
"AttemptStopRequest",
"204",
None,
),
(
"/internal/v1/attempts/{attempt_id}/staging/write-intent",
"AttemptStagingWriteIntentRequest",
"204",
None,
),
(
"/internal/v1/attempts/{attempt_id}/staging/object",
"AttemptStagingObjectRequest",
"204",
None,
),
(
"/internal/v1/attempts/{attempt_id}/staging/multipart",
"AttemptStagingMultipartRequest",
"204",
None,
),
(
"/internal/v1/attempts/{attempt_id}/resolved-source",
"AttemptResolvedSourceRequest",
"204",
None,
),
(
"/internal/v1/attempts/{attempt_id}/checkpoint",
"AttemptCheckpointRequest",
"204",
None,
),
(
"/internal/v1/attempts/{attempt_id}/progress",
"AttemptProgressRequest",
"204",
None,
),
(
"/internal/v1/attempts/{attempt_id}/outcome",
"AttemptOutcomeRequest",
"200",
Some("JobPhaseProjection"),
),
];
for (path, request, status, response) in routes {
assert_worker_route(&openapi, path, request, status, response);
}
}
fn assert_dto_and_schema<T>(component_schema: &Value, definition: &str, value: Value)
where
T: DeserializeOwned + Serialize,
{
let dto: T = serde_json::from_value(value.clone()).unwrap();
assert_eq!(serde_json::to_value(dto).unwrap(), value);
validate_openapi_value(component_schema, definition, &value);
}
fn validate_openapi_value(component_schema: &Value, definition: &str, value: &Value) {
let mut schema = component_schema.clone();
schema.as_object_mut().unwrap().insert(
"$ref".to_owned(),
Value::String(format!("#/$defs/{definition}")),
);
jsonschema::draft202012::options()
.should_validate_formats(true)
.build(&schema)
.unwrap()
.validate(value)
.unwrap();
}
fn merged(base: &Value, extension: Value) -> Value {
let mut value = base.clone();
value
.as_object_mut()
.unwrap()
.extend(extension.as_object().unwrap().clone());
value
}
fn assert_worker_route(
openapi: &Value,
path: &str,
request_schema: &str,
success_status: &str,
response_schema: Option<&str>,
) {
let operation = &openapi["paths"][path]["post"];
assert_eq!(
operation.pointer("/security/0/serviceAuth"),
Some(&serde_json::json!([]))
);
assert_eq!(
operation.pointer("/requestBody/content/application~1json/schema/$ref"),
Some(&Value::String(format!(
"#/components/schemas/{request_schema}"
)))
);
assert!(operation["responses"].get("409").is_some());
assert!(operation["responses"].get("default").is_some());
let success = &operation["responses"][success_status];
match response_schema {
Some(response_schema) => assert_eq!(
success.pointer("/content/application~1json/schema/$ref"),
Some(&Value::String(format!(
"#/components/schemas/{response_schema}"
)))
),
None => assert!(success.get("content").is_none()),
}
}
fn openapi_component_schema(openapi: &Value) -> Value {
let mut definitions = openapi
.pointer("/components/schemas")
.and_then(Value::as_object)
.unwrap()
.clone();
for schema in definitions.values_mut() {
rewrite_component_references(schema);
}
serde_json::json!({
"$schema": "https://json-schema.org/draft/2020-12/schema",
"$defs": definitions
})
}
fn rewrite_component_references(value: &mut Value) {
match value {
Value::Array(values) => {
for value in values {
rewrite_component_references(value);
}
}
Value::Object(object) => {
if let Some(Value::String(reference)) = object.get_mut("$ref")
&& let Some(name) = reference.strip_prefix("#/components/schemas/")
{
*reference = format!("#/$defs/{name}");
}
object.remove("discriminator");
for value in object.values_mut() {
rewrite_component_references(value);
}
}
_ => {}
}
}
fn validate_openapi_corpus(
component_schema: &Value,
corpus: &Value,
definition: &str,
corpus_key: &str,
) {
let mut schema = component_schema.clone();
schema.as_object_mut().unwrap().insert(
"$ref".to_owned(),
Value::String(format!("#/$defs/{definition}")),
);
let validator = jsonschema::draft202012::options()
.should_validate_formats(true)
.build(&schema)
.unwrap();
for value in corpus[corpus_key].as_array().unwrap() {
validator.validate(value).unwrap();
}
}
fn assert_local_references_resolve(document: &Value) {
let mut references = BTreeSet::new();
collect_local_references(document, &mut references);
assert!(!references.is_empty());
for reference in references {
let pointer = reference.strip_prefix('#').unwrap();
assert!(
document.pointer(pointer).is_some(),
"unresolved local reference {reference}"
);
}
}
fn collect_local_references(value: &Value, references: &mut BTreeSet<String>) {
match value {
Value::Array(values) => {
for value in values {
collect_local_references(value, references);
}
}
Value::Object(object) => {
if let Some(reference) = object.get("$ref").and_then(Value::as_str)
&& reference.starts_with("#/")
{
references.insert(reference.to_owned());
}
for value in object.values() {
collect_local_references(value, references);
}
}
_ => {}
}
}