Change test to check for correct SegmentationModel type.

spacetelescope · Aug 28, 2024 · 8fbf6b6 · 8fbf6b6
1 parent 2311211
commit 8fbf6b6
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/romancal/source_catalog/tests/test_source_catalog.py b/romancal/source_catalog/tests/test_source_catalog.py
@@ -14,6 +14,7 @@
     MosaicModel,
     MosaicSegmentationMapModel,
     MosaicSourceCatalogModel,
+    SegmentationMapModel,
     SourceCatalogModel,
 )
 from roman_datamodels.maker_utils import mk_level2_image, mk_level3_mosaic
@@ -461,7 +462,7 @@ def test_do_psf_photometry_column_names(tmp_path, image_model, fit_psf):
             ImageModel,
             {
                 "cat": SourceCatalogModel,
-                "segm": MosaicSegmentationMapModel,
+                "segm": SegmentationMapModel,
                 "sourcecatalog": ImageModel,
             },
         ),
@@ -474,7 +475,7 @@ def test_do_psf_photometry_column_names(tmp_path, image_model, fit_psf):
             SourceCatalogModel,
             {
                 "cat": SourceCatalogModel,
-                "segm": MosaicSegmentationMapModel,
+                "segm": SegmentationMapModel,
             },
         ),
         (
@@ -486,7 +487,7 @@ def test_do_psf_photometry_column_names(tmp_path, image_model, fit_psf):
             ImageModel,
             {
                 "cat": SourceCatalogModel,
-                "segm": MosaicSegmentationMapModel,
+                "segm": SegmentationMapModel,
             },
         ),
         (
@@ -498,7 +499,7 @@ def test_do_psf_photometry_column_names(tmp_path, image_model, fit_psf):
             SourceCatalogModel,
             {
                 "cat": SourceCatalogModel,
-                "segm": MosaicSegmentationMapModel,
+                "segm": SegmentationMapModel,
             },
         ),
     ),