Addressing PR comments

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

.coverage

@@ -4,6 +4,10 @@
 from pathlib import (
     Path,
 )
+from unittest.mock import (
+    MagicMock,
+    patch,
+)
 
 from deepmd.tf.entrypoints.change_bias import (
     change_bias,
@@ -89,6 +93,129 @@ def test_change_bias_user_defined_not_implemented(self):
             "User-defined bias setting is not yet implemented", str(cm.exception)
         )
 
+    def test_change_bias_successful_execution(self):
+        """Test successful bias changing execution path."""
+        # Create fake checkpoint directory with required files
+        fake_checkpoint_dir = self.temp_path / "checkpoint"
+        fake_checkpoint_dir.mkdir()
+        (fake_checkpoint_dir / "checkpoint").write_text("fake checkpoint content")
+        (fake_checkpoint_dir / "input.json").write_text('{"model": {}}')
+
+        # Create fake data system
+        fake_data_dir = self.temp_path / "data_system"
+        fake_data_dir.mkdir()
+        fake_set_dir = fake_data_dir / "set.000"
+        fake_set_dir.mkdir()
+
+        # Import the module properly
+        import sys
+
+        change_bias_module = sys.modules["deepmd.tf.entrypoints.change_bias"]
+
+        with (
+            patch.object(
+                change_bias_module, "expand_sys_str", return_value=[str(fake_data_dir)]
+            ),
+            patch.object(change_bias_module, "j_loader", return_value={"model": {}}),
+            patch.object(
+                change_bias_module, "update_deepmd_input", return_value={"model": {}}
+            ),
+            patch.object(change_bias_module, "normalize", return_value={"model": {}}),
+            patch.object(change_bias_module, "DeepmdDataSystem") as mock_data_system,
+            patch.object(change_bias_module, "DPTrainer") as mock_trainer_class,
+            patch.object(change_bias_module, "shutil"),
+        ):
+            # Mock the data system
+            mock_data_instance = MagicMock()
+            mock_data_instance.get_type_map.return_value = ["H", "O"]
+            mock_data_system.return_value = mock_data_instance
+
+            # Mock the trainer
+            mock_trainer_instance = MagicMock()
+            mock_model = MagicMock()
+            mock_model.get_type_map.return_value = ["H", "O"]
+            mock_trainer_instance.model = mock_model
+            mock_trainer_instance._change_energy_bias = MagicMock()
+            mock_trainer_instance.save_checkpoint = MagicMock()
+            mock_trainer_class.return_value = mock_trainer_instance
+
+            # Call change_bias function
+            change_bias(
+                INPUT=str(fake_checkpoint_dir),
+                mode="change",
+                system=str(fake_data_dir),
+                output=str(self.temp_path / "output"),
+            )
+
+            # Verify that the trainer's change_energy_bias was called
+            mock_trainer_instance._change_energy_bias.assert_called_once()
+
+    def test_change_bias_with_data_type_map(self):
+        """Test bias changing when data system has its own type_map."""
+        # Create fake checkpoint directory with required files
+        fake_checkpoint_dir = self.temp_path / "checkpoint"
+        fake_checkpoint_dir.mkdir()
+        (fake_checkpoint_dir / "checkpoint").write_text("fake checkpoint content")
+        (fake_checkpoint_dir / "input.json").write_text('{"model": {}}')
+
+        # Create fake data system
+        fake_data_dir = self.temp_path / "data_system"
+        fake_data_dir.mkdir()
+        fake_set_dir = fake_data_dir / "set.000"
+        fake_set_dir.mkdir()
+
+        # Import the module properly
+        import sys
+
+        change_bias_module = sys.modules["deepmd.tf.entrypoints.change_bias"]
+
+        with (
+            patch.object(
+                change_bias_module, "expand_sys_str", return_value=[str(fake_data_dir)]
+            ),
+            patch.object(change_bias_module, "j_loader", return_value={"model": {}}),
+            patch.object(
+                change_bias_module, "update_deepmd_input", return_value={"model": {}}
+            ),
+            patch.object(change_bias_module, "normalize", return_value={"model": {}}),
+            patch.object(change_bias_module, "DeepmdDataSystem") as mock_data_system,
+            patch.object(change_bias_module, "DPTrainer") as mock_trainer_class,
+            patch.object(change_bias_module, "shutil"),
+        ):
+            # Mock the data system with type_map
+            mock_data_instance = MagicMock()
+            mock_data_instance.get_type_map.return_value = [
+                "C",
+                "N",
+                "O",
+            ]  # Data has type_map
+            mock_data_system.return_value = mock_data_instance
+
+            # Mock the trainer
+            mock_trainer_instance = MagicMock()
+            mock_model = MagicMock()
+            mock_model.get_type_map.return_value = [
+                "H",
+                "O",
+            ]  # Model has different type_map
+            mock_trainer_instance.model = mock_model
+            mock_trainer_instance._change_energy_bias = MagicMock()
+            mock_trainer_instance.save_checkpoint = MagicMock()
+            mock_trainer_class.return_value = mock_trainer_instance
+
+            # Call change_bias function
+            change_bias(
+                INPUT=str(fake_checkpoint_dir),
+                mode="change",
+                system=str(fake_data_dir),
+            )
+
+            # Verify that data's type_map was used (not model's)
+            mock_trainer_instance._change_energy_bias.assert_called_once()
+            args, kwargs = mock_trainer_instance._change_energy_bias.call_args
+            # The third argument should be the type_map from data
+            self.assertEqual(args[2], ["C", "N", "O"])
+
 
 if __name__ == "__main__":
     unittest.main()