diff --git a/py/desidatamodel/test/test_check.py b/py/desidatamodel/test/test_check.py
index 08396150..4f988e13 100644
--- a/py/desidatamodel/test/test_check.py
+++ b/py/desidatamodel/test/test_check.py
@@ -3,11 +3,15 @@
 """Test desidatamodel.check functions
 """
 import os
-import sys
 from packaging import version
-import unittest
 from unittest.mock import patch
 
+skip_unit_warning = False
+try:
+    from desiutil.annotate import FITSUnitWarning
+except ImportError:
+    skip_unit_warning = True
+
 from .datamodeltestcase import DataModelTestCase, DM
 
 from .. import DataModelError
@@ -310,9 +314,10 @@ def test_extract_metadata_bad_keyword_unit(self):
             meta = model.extract_metadata(error=True)
         self.assertEqual(str(e.exception), erg_msg)
         model.hdumeta = None
-        meta = model.extract_metadata(error=False)
+        with self.assertWarns(FITSUnitWarning) as w:
+            meta = model.extract_metadata(error=False)
+        self.assertEqual(str(w.warning), erg_msg)
         self.assertLog(log, -1, "HDU 0 in {0} should have a more meaningful EXTNAME than 'PRIMARY'.".format(modelfile))
-        self.assertLog(log, -2, erg_msg)
 
     def test_extract_metadata_missing_keyword_unit(self):
         """Test reading metadata with missing units for header keywords.
@@ -347,8 +352,9 @@ def test_extract_metadata_bad_column_unit(self):
             meta = model.extract_metadata(error=True)
         self.assertEqual(str(e.exception), erg_msg)
         model.hdumeta = None
-        meta = model.extract_metadata(error=False)
-        self.assertLog(log, -1, erg_msg)
+        with self.assertWarns(FITSUnitWarning) as w:
+            meta = model.extract_metadata(error=False)
+        self.assertEqual(str(w.warning), erg_msg)
 
     def test_extract_metadata_missing_column_type(self):
         """Test reading metadata with missing FITS column types.
diff --git a/py/desidatamodel/test/test_model.py b/py/desidatamodel/test/test_model.py
index 18e14874..2d843128 100644
--- a/py/desidatamodel/test/test_model.py
+++ b/py/desidatamodel/test/test_model.py
@@ -3,7 +3,6 @@
 """Test data model files for validity.
 """
 import os
-import unittest
 
 from .datamodeltestcase import DataModelTestCase
 
diff --git a/py/desidatamodel/test/test_stub.py b/py/desidatamodel/test/test_stub.py
index aedab83c..a2a34edb 100644
--- a/py/desidatamodel/test/test_stub.py
+++ b/py/desidatamodel/test/test_stub.py
@@ -2,14 +2,17 @@
 # -*- coding: utf-8 -*-
 """Test desidatamodel.stub functions
 """
-import os
 import importlib.resources as ir
-import unittest
 from unittest.mock import patch, call
 from astropy.io import fits
 from astropy.io.fits.card import Undefined
 from collections import OrderedDict
 
+skip_unit_warning = False
+try:
+    from desiutil.annotate import FITSUnitWarning
+except ImportError:
+    skip_unit_warning = True
 from .datamodeltestcase import DataModelTestCase
 from .. import DataModelError
 from ..stub import (Stub, extrakey, file_size, fits_column_format,
@@ -381,9 +384,10 @@ def test_Stub_image_format(self):
         with self.assertRaises(ValueError) as e:
             i = stub.image_format(hdr)
         self.assertEqual(str(e.exception), erg_msg)
-        stub = Stub([sim_hdu(hdr)], error=False)
-        i = stub.image_format(hdr)
-        self.assertLog(log, 1, erg_msg)
+        with self.assertWarns(FITSUnitWarning) as w:
+            stub = Stub([sim_hdu(hdr)], error=False)
+            i = stub.image_format(hdr)
+        self.assertEqual(str(w.warning), erg_msg)
         self.assertEqual(i, 'Data: FITS image [BITPIX=128, 1000x1000]')
         hdr = sim_header()
         hdr['SIMPLE'] = True
@@ -394,7 +398,7 @@ def test_Stub_image_format(self):
         hdr['BUNIT'] = '10**-17 erg/(cm**2*s*Angstrom)'
         stub = Stub([sim_hdu(hdr)], error=False)
         i = stub.image_format(hdr)
-        self.assertLog(log, 4, "BUNIT   = '10**-17 erg/(cm**2*s*Angstrom)'")
+        self.assertLog(log, 2, "BUNIT   = '10**-17 erg/(cm**2*s*Angstrom)'")
         self.assertEqual(i, 'Data: FITS image [BITPIX=128, 1000x1000]')
         #
         # Check compressed HDU
diff --git a/py/desidatamodel/test/test_unit.py b/py/desidatamodel/test/test_unit.py
index a2da337a..2f66c8ca 100644
--- a/py/desidatamodel/test/test_unit.py
+++ b/py/desidatamodel/test/test_unit.py
@@ -2,8 +2,11 @@
 # -*- coding: utf-8 -*-
 """Test desidatamodel.unit functions
 """
-# import os
-import unittest
+skip_unit_warning = False
+try:
+    from desiutil.annotate import FITSUnitWarning
+except ImportError:
+    skip_unit_warning = True
 from .datamodeltestcase import DataModelTestCase
 from ..unit import validate_unit, log
 
@@ -18,12 +21,12 @@ def test_check_model(self):
         erg_msg = self.badUnitMessage('ergs')
         c = validate_unit('erg')
         self.assertIsNone(c)
-        c = validate_unit('ergs', error=False)
+        with self.assertWarns(FITSUnitWarning) as w:
+            c = validate_unit('ergs', error=False)
+        self.assertEqual(str(w.warning), erg_msg)
         self.assertIsNone(c)
         c = validate_unit('nanomaggies', error=True)
         self.assertEqual(c, "'nanomaggies'")
-        self.assertLog(log, -1, erg_msg)
         with self.assertRaises(ValueError) as e:
             c = validate_unit('ergs', error=True)
         self.assertEqual(str(e.exception), erg_msg)
-        self.assertLog(log, -1, erg_msg)