From e90dc84df945e1792812498e777126d9dbb67ff6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Agust=C3=ADn=20Borgna?= <agustin.borgna@quantinuum.com>
Date: Wed, 14 May 2025 14:10:54 +0100
Subject: [PATCH] test: Check hugr json serializations against the schema
 (again)

---
 .github/workflows/ci-rs.yml            |  1 +
 hugr-core/src/envelope/package_json.rs | 11 +++++++++++
 hugr-core/src/hugr/serialize/test.rs   | 12 ++++++++++++
 justfile                               |  4 ++--
 4 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/ci-rs.yml b/.github/workflows/ci-rs.yml
index c5f6609b8..c6814fc60 100644
--- a/.github/workflows/ci-rs.yml
+++ b/.github/workflows/ci-rs.yml
@@ -19,6 +19,7 @@ env:
   CI: true # insta snapshots behave differently on ci
   SCCACHE_GHA_ENABLED: "true"
   RUSTC_WRAPPER: "sccache"
+  HUGR_TEST_SCHEMA: "1"
   # different strings for install action and feature name
   # adapted from https://github.com/TheDan64/inkwell/blob/master/.github/workflows/test.yml
   LLVM_VERSION: "14.0"
diff --git a/hugr-core/src/envelope/package_json.rs b/hugr-core/src/envelope/package_json.rs
index 6a2b62b75..3bc7fe748 100644
--- a/hugr-core/src/envelope/package_json.rs
+++ b/hugr-core/src/envelope/package_json.rs
@@ -51,6 +51,17 @@ pub(super) fn to_json_writer<'h>(
         modules: hugrs.into_iter().map(HugrSer).collect(),
         extensions: extensions.iter().map(std::convert::AsRef::as_ref).collect(),
     };
+
+    // Validate the hugr serializations against the schema.
+    #[cfg(all(test, not(miri)))]
+    if std::env::var("HUGR_TEST_SCHEMA").is_ok_and(|x| !x.is_empty()) {
+        use crate::hugr::serialize::test::check_hugr_serialization_schema;
+
+        for hugr in &pkg_ser.modules {
+            check_hugr_serialization_schema(hugr.0);
+        }
+    }
+
     serde_json::to_writer(writer, &pkg_ser)?;
     Ok(())
 }
diff --git a/hugr-core/src/hugr/serialize/test.rs b/hugr-core/src/hugr/serialize/test.rs
index 2b500ed03..5b224b168 100644
--- a/hugr-core/src/hugr/serialize/test.rs
+++ b/hugr-core/src/hugr/serialize/test.rs
@@ -175,6 +175,18 @@ fn ser_roundtrip_check_schema<TSer: Serialize, TDeser: serde::de::DeserializeOwn
     serde_json::from_value(val).unwrap()
 }
 
+/// Serialize a Hugr and check that it is valid against the schema.
+///
+/// # Panics
+///
+/// Panics if the serialization fails or if the schema validation fails.
+pub(crate) fn check_hugr_serialization_schema(hugr: &Hugr) {
+    let schemas = get_schemas(true);
+    let hugr_ser = HugrSer(hugr);
+    let val = serde_json::to_value(hugr_ser).unwrap();
+    NamedSchema::check_schemas(&val, schemas);
+}
+
 /// Serialize and deserialize a HUGR, and check that the result is the same as the original.
 ///
 /// If `check_schema` is true, checks the serialized json against the in-tree schema.
diff --git a/justfile b/justfile
index 87771f980..a092d45bf 100644
--- a/justfile
+++ b/justfile
@@ -10,7 +10,7 @@ setup:
 
 # Run the pre-commit checks.
 check:
-    uv run pre-commit run --all-files
+    HUGR_TEST_SCHEMA=1 uv run pre-commit run --all-files
 
 # Run all the tests.
 test: test-rust test-python
@@ -20,7 +20,7 @@ test-rust *TEST_ARGS:
     @# built into a binary build (without using `maturin`)
     @#
     @# This feature list should be kept in sync with the `hugr-py/pyproject.toml`
-    cargo test \
+    HUGR_TEST_SCHEMA=1 cargo test \
         --workspace \
         --exclude 'hugr-py' \
         --features 'hugr/declarative hugr/llvm hugr/llvm-test hugr/zstd' {{TEST_ARGS}}