def test_deepspeed_load_checkpoint_client_state_missing(tmp_path):
    """Test that the DeepSpeed strategy raises a custom error when client state couldn't be loaded by DeepSpeed."""
    from deepspeed import DeepSpeedEngine

    strategy = DeepSpeedStrategy()
    optimizer = Mock()
    model = Mock(spec=DeepSpeedEngine, optimizer=optimizer)
    model.modules.return_value = [model]

    # If the DeepSpeed engine fails to load the checkpoint file (e.g., file not found), it prints a warning and
    # returns None from its function call
    model.load_checkpoint.return_value = [None, None]

    # Check for our custom user error
    with pytest.raises(FileNotFoundError, match="The provided path is not a valid DeepSpeed checkpoint"):
        strategy.load_checkpoint(path=tmp_path, state={"model": model, "optimizer": optimizer, "test": "data"})
Test that the DeepSpeed strategy raises a custom error when client state couldn't be loaded by DeepSpeed.
tests/tests_fabric/strategies/test_deepspeed.py