from __future__ import print_function, division, absolute_import

import pytest

from distributed.core import rpc
from distributed.utils_test import gen_cluster
from distributed.utils_comm import pack_data, gather_from_workers


def test_pack_data():
    data = {"x": 1}
    assert pack_data(("x", "y"), data) == (1, "y")
    assert pack_data({"a": "x", "b": "y"}, data) == {"a": 1, "b": "y"}
    assert pack_data({"a": ["x"], "b": "y"}, data) == {"a": [1], "b": "y"}


@pytest.mark.xfail(reason="rpc now needs to be a connection pool")
@gen_cluster(client=True)
def test_gather_from_workers_permissive(c, s, a, b):
    x = yield c.scatter({"x": 1}, workers=a.address)

    data, missing, bad_workers = yield gather_from_workers(
        {"x": [a.address], "y": [b.address]}, rpc=rpc
    )

    assert data == {"x": 1}
    assert list(missing) == ["y"]