From 2c76bc25d14e87f3b19d4cf645d817d55dd6b808 Mon Sep 17 00:00:00 2001
From: Caleb Schilly <cwschilly@gmail.com>
Date: Fri, 13 Sep 2024 12:32:38 -0400
Subject: [PATCH] #2201: fix bug in schema; require collection_id for
 migratable objects

---
 scripts/JSON_data_files_validator.py |  9 +++++----
 scripts/LBDatafile_schema.py         | 17 ++++++++++++-----
 2 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/scripts/JSON_data_files_validator.py b/scripts/JSON_data_files_validator.py
index 72a5c51c38..a14c083fc6 100644
--- a/scripts/JSON_data_files_validator.py
+++ b/scripts/JSON_data_files_validator.py
@@ -434,13 +434,14 @@ def validate_comm_links(all_jsons):
 
             for data in all_jsons:
                 tasks = data["phases"][n]["tasks"]
-                id_key = "id" if "id" in tasks[0]["entity"] else "seq_id"
-                task_ids.update({int(task["entity"][id_key]) for task in tasks})
+                task_ids.update(
+                    {int(task["entity"].get("id", task["entity"].get("seq_id"))) for task in tasks}
+                )
 
                 if data["phases"][n].get("communications") is not None:
                     comms = data["phases"][n]["communications"]
-                    comm_ids.update({int(comm["from"][id_key]) for comm in comms})
-                    comm_ids.update({int(comm["to"][id_key]) for comm in comms})
+                    comm_ids.update({int(comm["from"].get("id", comm["from"].get("seq_id"))) for comm in comms})
+                    comm_ids.update({int(comm["to"].get("id", comm["to"].get("seq_id"))) for comm in comms})
 
             if not comm_ids.issubset(task_ids):
                 logging.error(
diff --git a/scripts/LBDatafile_schema.py b/scripts/LBDatafile_schema.py
index 743fff574e..d11b64fa0e 100644
--- a/scripts/LBDatafile_schema.py
+++ b/scripts/LBDatafile_schema.py
@@ -1,9 +1,16 @@
 from schema import And, Optional, Schema
 
-def validate_id_and_seq_id(field):
-    """Ensure that either seq_id or id is provided."""
+def validate_ids(field):
+    """
+    Ensure that 1) either seq_id or id is provided,
+    and 2) if an object is migratable, collection_id has been set.
+    """
     if 'seq_id' not in field and 'id' not in field:
         raise ValueError('Either id (bit-encoded) or seq_id must be provided.')
+
+    if field['migratable'] and 'collection_id' not in field:
+        raise ValueError('If an entity is migratable, it must have a collection_id')
+
     return field
 
 LBDatafile_schema = Schema(
@@ -45,7 +52,7 @@ def validate_id_and_seq_id(field):
                             'type': str,
                             'migratable': bool,
                             Optional('objgroup_id'): int
-                        }, validate_id_and_seq_id),
+                        }, validate_ids),
                         'node': int,
                         'resource': str,
                         Optional('subphases'): [
@@ -71,7 +78,7 @@ def validate_id_and_seq_id(field):
                             Optional('migratable'): bool,
                             Optional('index'): [int],
                             Optional('objgroup_id'): int,
-                        }, validate_id_and_seq_id),
+                        }, validate_ids),
                         'messages': int,
                         'from': And({
                             'type': str,
@@ -82,7 +89,7 @@ def validate_id_and_seq_id(field):
                             Optional('migratable'): bool,
                             Optional('index'): [int],
                             Optional('objgroup_id'): int,
-                        }, validate_id_and_seq_id),
+                        }, validate_ids),
                         'bytes': float
                     }
                 ],