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

use std::{collections::BTreeSet, fs};

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");
}

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);
            }
        }
        _ => {}
    }
}