from __future__ import print_function, division, absolute_import

from distributed.protocol.utils import merge_frames, pack_frames, unpack_frames
from distributed.utils import ensure_bytes


def test_merge_frames():
    result = merge_frames({'lengths': [3, 4]}, [b'12', b'34', b'567'])
    expected = [b'123', b'4567']

    assert list(map(ensure_bytes, result)) == expected

    b = b'123'
    assert merge_frames({'lengths': [3]}, [b])[0] is b

    L = [b'123', b'456']
    assert merge_frames({'lengths': [3, 3]}, L) is L


def test_pack_frames():
    frames = [b'123', b'asdf']
    b = pack_frames(frames)
    assert isinstance(b, bytes)
    frames2 = unpack_frames(b)

    assert frames == frames2