def test_init_local_comm(monkeypatch):
    """Test that it can determine the node rank and local rank based on the hostnames of all participating nodes."""
    # pretend mpi4py is available
    monkeypatch.setattr(lightning.fabric.plugins.environments.mpi, "_MPI4PY_AVAILABLE", True)
    mpi4py_mock = MagicMock()
    hostname_mock = MagicMock()

    mpi4py_mock.MPI.COMM_WORLD.Get_size.return_value = 4
    with mock.patch.dict("sys.modules", {"mpi4py": mpi4py_mock}), mock.patch("socket.gethostname", hostname_mock):
        env = MPIEnvironment()

        hostname_mock.return_value = "host1"
        env._comm_world.gather.return_value = ["host1", "host2"]
        env._comm_world.bcast.return_value = ["host1", "host2"]
        assert env.node_rank() == 0

        env._node_rank = None
        hostname_mock.return_value = "host2"
        env._comm_world.gather.return_value = None
        env._comm_world.bcast.return_value = ["host1", "host2"]
        assert env.node_rank() == 1

        assert env._comm_local is not None
        env._comm_local.Get_rank.return_value = 33
        assert env.local_rank() == 33
Test that it can determine the node rank and local rank based on the hostnames of all participating nodes.
tests/tests_fabric/plugins/environments/test_mpi.py