From d1a54657c2d35cf25856f94cc195d05055a45fec Mon Sep 17 00:00:00 2001
From: Daniel King <daniel@mosaicml.com>
Date: Fri, 26 Jul 2024 12:16:04 -0700
Subject: [PATCH] add test

---
 tests/a_scripts/inference/test_convert_composer_to_hf.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tests/a_scripts/inference/test_convert_composer_to_hf.py b/tests/a_scripts/inference/test_convert_composer_to_hf.py
index ffdb09ca98..9eb214e83d 100644
--- a/tests/a_scripts/inference/test_convert_composer_to_hf.py
+++ b/tests/a_scripts/inference/test_convert_composer_to_hf.py
@@ -388,6 +388,9 @@ def test_huggingface_conversion_callback_interval(
     checkpointer_callback.transform_model_pre_registration = MagicMock(
         wraps=checkpointer_callback.transform_model_pre_registration,
     )
+    checkpointer_callback.pre_register_edit = MagicMock(
+        wraps=checkpointer_callback.pre_register_edit,
+    )
     trainer = Trainer(
         model=original_model,
         device='gpu',
@@ -413,9 +416,11 @@ def test_huggingface_conversion_callback_interval(
             metadata={},
         )
         assert checkpointer_callback.transform_model_pre_registration.call_count == 1
+        assert checkpointer_callback.pre_register_edit.call_count == 1
         assert mlflow_logger_mock.register_model_with_run_id.call_count == 1
     else:
         assert checkpointer_callback.transform_model_pre_registration.call_count == 0
+        assert checkpointer_callback.pre_register_edit.call_count == 0
         assert mlflow_logger_mock.save_model.call_count == 0
         assert mlflow_logger_mock.register_model_with_run_id.call_count == 0