from psutil import virtual_memory


def mock_cluster(n_workers=1,
                 threads_per_worker=1,
                 diagnostics_port=8787,
                 memory_limit=None,
                 **dask_kwarg):
    return (n_workers, threads_per_worker, diagnostics_port, memory_limit)


class MockClient():
    def __init__(self, cluster):
        self.cluster = cluster

    def scheduler_info(self):
        return {'workers': {'worker 1': {'memory_limit': virtual_memory().total}}}


def get_mock_client_cluster():
    return MockClient, mock_cluster
