def test_step_method_redirection():
    """Test that the FabricModule redirects methods marked as 'forward methods' through forward to avoid bypassing the
    DDP/FSDP wrappers."""

    class DDP(torch.nn.Module):
        def __init__(self, module):
            super().__init__()
            self.module = module

        def forward(self, *args, **kwargs):
            return self.module(*args, **kwargs)

    class LightningModule(torch.nn.Module):
        def forward(self):
            return "forward_return"

        def training_step(self, arg, kwarg=None):
            assert self() == "forward_return"
            assert arg == "train_arg"
            assert kwarg == "train_kwarg"
            return "training_step_return"

        def marked_method(self, arg, kwarg=None):
            assert self() == "forward_return"
            assert arg == "marked_arg"
            assert kwarg == "marked_kwarg"
            return "marked_method_return"

        def normal_method(self):
            pass

    strategy = Mock()
    strategy.precision = Mock(wraps=Precision())
    original_module = LightningModule()
    forward_module = DDP(original_module)
    fabric_module = _FabricModule(forward_module=forward_module, strategy=strategy, original_module=original_module)

    # Regular methods on the original_module are visible and identical on the fabric_module ...
    assert fabric_module.normal_method.__wrapped__ == original_module.normal_method

    # ... but special methods like training_step get redirected to the forward_module
    assert fabric_module.training_step.__name__ == "call_forward_module"
    assert fabric_module.validation_step.__name__ == "call_forward_module"
    assert fabric_module.test_step.__name__ == "call_forward_module"
    assert fabric_module.predict_step.__name__ == "call_forward_module"

    with pytest.raises(AttributeError, match="has no attribute 'predict_step'"):
        # A special method that does not exist will raise its AttributeError when being called
        fabric_module.predict_step()

    # The forward method on the original module remains untouched
    assert original_module.forward.__name__ == "forward"

    # The special methods get redirected correctly to produce the expected output
    strategy.precision.forward_context.reset_mock()
    assert fabric_module.training_step("train_arg", kwarg="train_kwarg") == "training_step_return"
    assert fabric_module.training_step("train_arg", kwarg="train_kwarg") == "training_step_return"  # call 2nd time
    assert strategy.precision.forward_context.call_count == 2

    # Other methods must be marked explicitly to be redirected
    strategy.precision.forward_context.reset_mock()
    with pytest.raises(RuntimeError, match="You are calling the method .* from outside the model"):
        fabric_module.marked_method("marked_arg", kwarg="marked_kwarg")
    fabric_module.mark_forward_method("marked_method")
    assert fabric_module.marked_method("marked_arg", kwarg="marked_kwarg") == "marked_method_return"
    strategy.precision.forward_context.assert_called_once()

    # The forward method remains untouched/unpatched after the special methods have been called
    assert original_module.forward.__name__ == "forward"

    # Special case: forward_module == original_module -> no special treatment applied
    fabric_module = _FabricModule(forward_module=original_module, strategy=Mock(), original_module=original_module)
    assert fabric_module.training_step == original_module.training_step
    assert fabric_module.marked_method == original_module.marked_method
Test that the FabricModule redirects methods marked as 'forward methods' through forward to avoid bypassing the
    DDP/FSDP wrappers.
tests/tests_fabric/test_wrappers.py