Repository: dask/dask
File: dask/array/tests/test_array_core.py (Python)

Code before:
        assert (ar == 2).all()


def test_store_compute_false():
    d = da.ones((4, 4), chunks=(2, 2))
    a, b = d + 1, d + 2

    at = np.zeros(shape=(4, 4))
    bt = np.zeros(shape=(4, 4))

    v = store([a, b], [at, bt], compute=False)
    assert isinstance(v, Delayed)

    # You need a well-formed HighLevelgraph for e.g. dask.graph_manipulation.bind
    for layer in v.__dask_layers__():
        assert layer in v.dask.layers

    assert (at == 0).all() and (bt == 0).all()
    assert all([ev is None for ev in v.compute()])
    assert (at == 2).all() and (bt == 3).all()

    at = np.zeros(shape=(4, 4))
    bt = np.zeros(shape=(4, 4))

    for load_stored in [None, False, True]:
        dat, dbt = store(
            [a, b], [at, bt], compute=False, return_stored=True, load_stored=load_stored
        )
        assert isinstance(dat, Array) and isinstance(dbt, Array)
        assert (at == 0).all() and (bt == 0).all()
        assert (dat.compute() == at).all() and (dbt.compute() == bt).all()
        assert (at == 2).all() and (bt == 3).all()


def test_store_nocompute_regions():
    x = da.ones(10, chunks=1)
    y = np.zeros((2, 10))
    d1 = da.store(x, y, regions=(0,), compute=False)
    d2 = da.store(x, y, regions=(1,), compute=False)
    assert d1.key != d2.key


class ThreadSafetyError(Exception):
    pass


class NonthreadSafeStore:
    def __init__(self):
        self.in_use = False

    def __setitem__(self, key, value):

Review comment (suggestion):
is there a way to assert that `compute` is not called in here?

Xarray has this: https://github.com/pydata/xarray/blob/65b058144ab16bd03e4b783955031e27b28d1311/xarray/tests/__init__.py#L225-L230

Code after:
        assert (ar == 2).all()


def test_store_compute_false():
    d = da.ones((4, 4), chunks=(2, 2))
    a, b = d + 1, d + 2

    at = np.zeros(shape=(4, 4))
    bt = np.zeros(shape=(4, 4))

    v = store([a, b], [at, bt], compute=False)
    assert isinstance(v, Delayed)

    # You need a well-formed HighLevelgraph for e.g. dask.graph_manipulation.bind
    for layer in v.__dask_layers__():
        assert layer in v.dask.layers

    assert (at == 0).all() and (bt == 0).all()
    assert all([ev is None for ev in v.compute()])
    assert (at == 2).all() and (bt == 3).all()

    at = np.zeros(shape=(4, 4))
    bt = np.zeros(shape=(4, 4))
    dat, dbt = store([a, b], [at, bt], compute=False, return_stored=True)
    assert isinstance(dat, Array) and isinstance(dbt, Array)
    assert (at == 0).all() and (bt == 0).all()
    assert (dat.compute() == at).all() and (dbt.compute() == bt).all()
    assert (at == 2).all() and (bt == 3).all()

    at = np.zeros(shape=(4, 4))
    bt = np.zeros(shape=(4, 4))
    dat, dbt = store(
        [a, b], [at, bt], compute=False, return_stored=True, load_stored=False
    )
    assert isinstance(dat, Array) and isinstance(dbt, Array)
    assert (at == 0).all() and (bt == 0).all()
    dask.compute(dat, dbt)
    assert (at == 2).all() and (bt == 3).all()


def test_store_nocompute_regions():
    x = da.ones(10, chunks=1)
    y = np.zeros((2, 10))
    d1 = da.store(x, y, regions=(0,), compute=False)
    d2 = da.store(x, y, regions=(1,), compute=False)
    assert d1.key != d2.key


class ThreadSafetyError(Exception):
    pass