From 633dcea3f8f283b41f138a10af75347669dc36cd Mon Sep 17 00:00:00 2001
From: William Jamieson <william@jamieson.phd>
Date: Fri, 29 Nov 2024 10:53:27 -0500
Subject: [PATCH] Add more complete tests of roman and jwst datamodels

---
 tests/test_jwst.py      | 31 +++++++++++++++++++++++++++++++
 tests/test_protocols.py | 17 -----------------
 tests/test_roman.py     | 35 +++++++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+), 17 deletions(-)
 create mode 100644 tests/test_jwst.py
 create mode 100644 tests/test_roman.py

diff --git a/tests/test_jwst.py b/tests/test_jwst.py
new file mode 100644
index 0000000..4bb300b
--- /dev/null
+++ b/tests/test_jwst.py
@@ -0,0 +1,31 @@
+"""
+Integration tests with JWST pipeline
+"""
+
+from inspect import getmembers, isclass
+
+import pytest
+
+from stpipe.protocols import DataModel
+
+datamodels = pytest.importorskip("stdatamodels.jwst.datamodels")
+
+
+def test_jwst_datamodel():
+    """Smoke test to ensure the JWST datamodels work with the DataModel protocol."""
+    jwst_datamodel = pytest.importorskip("stdatamodels.jwst.datamodels")
+    image_model = jwst_datamodel.ImageModel()
+    assert isinstance(image_model, DataModel)
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        model[1]
+        for model in getmembers(datamodels, isclass)
+        if issubclass(model[1], datamodels.JwstDataModel)
+    ],
+)
+def test_datamodel(model):
+    """Test that all JWST datamodels work with the DataModel protocol."""
+    assert isinstance(model(), DataModel)
diff --git a/tests/test_protocols.py b/tests/test_protocols.py
index 2acb071..eb4239a 100644
--- a/tests/test_protocols.py
+++ b/tests/test_protocols.py
@@ -2,26 +2,9 @@
 Test that the DataModel interface of the protocol works properly
 """
 
-import pytest
-
 from stpipe.protocols import DataModel
 
 
-def test_roman_datamodel():
-    roman_datamodels = pytest.importorskip("roman_datamodels.datamodels")
-    from roman_datamodels.maker_utils import mk_level2_image
-
-    roman_image_tree = mk_level2_image()
-    image_model = roman_datamodels.ImageModel(roman_image_tree)
-    assert isinstance(image_model, DataModel)
-
-
-def test_jwst_datamodel():
-    jwst_datamodel = pytest.importorskip("stdatamodels.jwst.datamodels")
-    image_model = jwst_datamodel.ImageModel()
-    assert isinstance(image_model, DataModel)
-
-
 class GoodDataModel:
     def __init__(self):
         pass
diff --git a/tests/test_roman.py b/tests/test_roman.py
new file mode 100644
index 0000000..7e7c31f
--- /dev/null
+++ b/tests/test_roman.py
@@ -0,0 +1,35 @@
+"""
+Integration tests with Roman pipeline
+"""
+
+from inspect import getmembers, isclass
+
+import pytest
+
+from stpipe.protocols import DataModel
+
+datamodels = pytest.importorskip("roman_datamodels.datamodels")
+
+
+def test_roman_datamodel():
+    """Smoke test to ensure the Roman datamodels work with the DataModel protocol."""
+    roman_datamodels = pytest.importorskip("roman_datamodels.datamodels")
+    from roman_datamodels.maker_utils import mk_level2_image
+
+    roman_image_tree = mk_level2_image()
+    image_model = roman_datamodels.ImageModel(roman_image_tree)
+    assert isinstance(image_model, DataModel)
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        model[1]
+        for model in getmembers(datamodels, isclass)
+        if model[1] != datamodels.DataModel
+        and issubclass(model[1], datamodels.DataModel)
+    ],
+)
+def test_datamodel(model):
+    """Test that all Roman datamodels work with the DataModel protocol."""
+    assert isinstance(model(), DataModel)