wan_checkpointer_test.py corrected

Author: prishajain1

src/maxdiffusion/tests/wan_checkpointer_test.py

@@ -408,8 +408,12 @@ def test_load_checkpoint_with_optimizer_in_low_noise(self, mock_wan_pipeline_i2v
     checkpointer = WanCheckpointerI2V_2_2(config=self.config)
     pipeline, opt_state, step = checkpointer.load_checkpoint(step=1)
 
+    mock_manager.restore.assert_called_once()
+    mock_wan_pipeline_i2v_2p2.from_checkpoint.asset_called_once_with(self.config, restored_mock)
+    self.assertEqual(pipeline, mock_pipeline_instance)
     self.assertIsNotNone(opt_state)
     self.assertEqual(opt_state["learning_rate"], 0.001)
+    self.assertEqual(step,1)
 
   @patch("maxdiffusion.checkpointing.wan_checkpointer.create_orbax_checkpoint_manager")
   @patch("maxdiffusion.checkpointing.wan_checkpointer_i2v_2p2.WanPipelineI2V_2_2")
@@ -436,8 +440,12 @@ def test_load_checkpoint_with_optimizer_in_high_noise(self, mock_wan_pipeline_i2
     checkpointer = WanCheckpointerI2V_2_2(config=self.config)
     pipeline, opt_state, step = checkpointer.load_checkpoint(step=1)
 
+    mock_manager.restore.assert_called_once()
+    mock_wan_pipeline_i2v_2p2.from_checkpoint.assert_called_once_with(self.config, restored_mock)
+    self.assertEqual(pipeline, mock_pipeline_instance)
     self.assertIsNotNone(opt_state)
     self.assertEqual(opt_state["learning_rate"], 0.002)
+    self.assertEqual(step, 1)
 
 class WanCheckpointerEdgeCasesTest(unittest.TestCase):
   """Tests for edge cases and error handling."""