diff --git a/src/maxdiffusion/tests/models_import_test.py b/src/maxdiffusion/tests/models_import_test.py
index 3bd898cde..45610ed0e 100644
--- a/src/maxdiffusion/tests/models_import_test.py
+++ b/src/maxdiffusion/tests/models_import_test.py
@@ -76,6 +76,7 @@ def test_mldiagnostics_import_and_usage(self):
     try:
       from google_cloud_mldiagnostics import machinelearning_run, xprof
       self.assertIsNotNone(machinelearning_run)
+      self.assertIsNotNone(xprof)
     except ImportError:
       pass