def test_rank_zero_first(monkeypatch):
    """Test that rank 0 completes first before all other processes can execute under `.rank_zero_first()`."""

    def record_calls_for_rank(rank):
        call_order = []

        fabric = Fabric()
        fabric._strategy = Mock(global_rank=rank)
        barrier_mock = MagicMock(side_effect=lambda *_: call_order.append("barrier"))
        monkeypatch.setattr(lightning.fabric.utilities.distributed._InfiniteBarrier, "__call__", barrier_mock)
        target = Mock(run=Mock(side_effect=lambda *_: call_order.append("run")))

        with fabric.rank_zero_first():
            target.run()

        return call_order

    assert record_calls_for_rank(0) == ["run", "barrier"]
    assert record_calls_for_rank(1) == ["barrier", "run"]
Test that rank 0 completes first before all other processes can execute under `.rank_zero_first()`.
tests/tests_fabric/test_fabric.py