def test_module_sharding_context():
    """Test that the sharding context manager gets applied when the strategy supports it and is a no-op otherwise."""
    fabric = Fabric()
    fabric._launched = True
    fabric._strategy = MagicMock(spec=DDPStrategy, module_sharded_context=Mock())
    with pytest.warns(DeprecationWarning, match="sharded_model"), fabric.sharded_model():
        pass
    fabric._strategy.module_sharded_context.assert_not_called()

    fabric._strategy = MagicMock(spec=_Sharded)
    with pytest.warns(DeprecationWarning, match="sharded_model"), fabric.sharded_model():
        pass
    fabric._strategy.module_sharded_context.assert_called_once()
Test that the sharding context manager gets applied when the strategy supports it and is a no-op otherwise.
tests/tests_fabric/test_fabric.py