<commit_before>import os

import pytest
pytest.importorskip('zmq')
pytest.importorskip('dill')

from dask.distributed import Worker, Scheduler, Client
from dask.utils import raises
from contextlib import contextmanager
from operator import add
from time import sleep
from multiprocessing.pool import ThreadPool
from toolz import partial

def inc(x):
    return x + 1

@contextmanager
def scheduler_and_workers(n=2):
    s = Scheduler()
    workers = [Worker(s.address_to_workers) for i in range(n)]
    while len(s.workers) < n:
        sleep(0.01)
    try:
        yield s, workers
    finally:
        s.close()
        for w in workers:
            w.close()


def test_get():
    with scheduler_and_workers() as (s, (a, b)):
        c = Client(s.address_to_clients)

        dsk = {'x': 1, 'y': (add, 'x', 'x'), 'z': (inc, 'y')}
        keys = ['y', 'z']

        assert c.get(dsk, keys) == [2, 3]
        c.close()


def test_status():
    with scheduler_and_workers() as (s, (a, b)):
        c = Client(s.address_to_clients)

        assert c.scheduler_status() == 'OK'
        c.close()


def test_get_with_dill():
    with scheduler_and_workers() as (s, (a, b)):
        c = Client(s.address_to_clients)<commit_msg>test micro-optimization.<commit_after>import os

import pytest
pytest.importorskip('zmq')
pytest.importorskip('dill')

from dask.distributed import Worker, Scheduler, Client
from dask.utils import raises
from contextlib import contextmanager
from operator import add
from time import sleep
from multiprocessing.pool import ThreadPool
from toolz import partial

def inc(x):
    return x + 1

@contextmanager
def scheduler_and_workers(n=2):
    s = Scheduler()
    workers = [Worker(s.address_to_workers) for i in range(n)]
    while len(s.workers) < n:
        sleep(1e-6)
    try:
        yield s, workers
    finally:
        s.close()
        for w in workers:
            w.close()


def test_get():
    with scheduler_and_workers() as (s, (a, b)):
        c = Client(s.address_to_clients)

        dsk = {'x': 1, 'y': (add, 'x', 'x'), 'z': (inc, 'y')}
        keys = ['y', 'z']

        assert c.get(dsk, keys) == [2, 3]
        c.close()


def test_status():
    with scheduler_and_workers() as (s, (a, b)):
        c = Client(s.address_to_clients)

        assert c.scheduler_status() == 'OK'
        c.close()


def test_get_with_dill():
    with scheduler_and_workers() as (s, (a, b)):
        c = Client(s.address_to_clients)
{'fraction_non_alphanumeric': 0.13186813186813187, 'fraction_numerical': 0.006105006105006105, 'mean_word_length': 2.9140127388535033, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6871795', 'n_tokens_mistral': 873, 'n_tokens_neox': 854, 'n_words': 240}