always pair astropy and asdf unit/quantity tags

tests/test_schemas.py

@@ -1,13 +1,16 @@
 import re
 
 import asdf
+import pytest
 
 ALLOWED_REFS = (
     r"^baseframe-[0-9.]+$",
     r"^frames/baseframe-[0-9.]+$",
     r"^#.*$",
 )
 
+UNIT_TAGS = {"tag:stsci.edu:asdf/unit/unit-1.*", "tag:astropy.org:astropy/units/unit-1.*"}
+
 
 def test_only_known_refs(latest_schema):
     """Latest schemas should only contain specific refs"""
@@ -29,3 +32,38 @@ def test_wildcard_tags(latest_schema):
             pattern = node["tag"]
             if "*" not in pattern:
                 assert False, f"tag pattern missing wildcard: {pattern}"
+
+
+@pytest.mark.parametrize("tag_set", (UNIT_TAGS,))
+def test_tags_in_allof(latest_schema, tag_set):
+    """
+    Test that some tags (unit) where the
+    tag used depends on the value are always referenced in an
+    allof containing all tags.
+    """
+    # we walk_and_modify here to use postorder
+
+    def callback(node):
+        if not isinstance(node, dict):
+            return node
+        if "anyOf" in node:
+            # check if this anyof includes a tag in the set
+            seen = set()
+            for sub in node["anyOf"]:
+                if not isinstance(sub, dict):
+                    continue
+                if "tag" not in sub:
+                    continue
+                if sub["tag"] in tag_set:
+                    seen.add(sub["tag"])
+
+            if seen:
+                # if a tag was found, check both were found
+                assert seen == tag_set, f"anyOf {node} missing: {tag_set - seen}"
+                # remove the anyof so the code below can check for tags not in the anyof
+                return {k: v for k, v in node.items() if k != "anyOf"}
+        if tag := node.get("tag"):
+            assert tag not in tag_set, f"tag {tag} must be in an anyOf with: {tag_set}"
+        return node
+
+    asdf.treeutil.walk_and_modify(latest_schema, callback, postorder=False)