import pandas as pd
import pytest

pa = pytest.importorskip('pyarrow')

from distributed.utils_test import gen_cluster
from distributed.protocol import deserialize, serialize


df = pd.DataFrame({'A': list('abc'), 'B': [1,2,3]})
tbl = pa.Table.from_pandas(df, preserve_index=False)
batch = pa.RecordBatch.from_pandas(df, preserve_index=False)


@pytest.mark.parametrize('obj', [batch, tbl], ids=["RecordBatch", "Table"])
def test_roundtrip(obj):
    # Test that the serialize/deserialize functions actually
    # work independent of distributed
    header, frames = serialize(obj)
    new_obj = deserialize(header, frames)
    assert obj.equals(new_obj)


def echo(arg):
    return arg


@pytest.mark.parametrize('obj', [batch, tbl], ids=["RecordBatch", "Table"])
def test_scatter(obj):
    @gen_cluster(client=True)
    def run_test(client, scheduler, worker1, worker2):
        obj_fut = yield client.scatter(obj)
        fut = client.submit(echo, obj_fut)
        result = yield fut
        assert obj.equals(result)
    run_test()