from pymemcache.client.rendezvous import RendezvousHash
import pytest


@pytest.mark.unit()
def test_init_no_options():
    rendezvous = RendezvousHash()
    assert 0 == len(rendezvous.nodes)
    assert 1361238019 == rendezvous.hash_function("6666")


@pytest.mark.unit()
def test_init():
    nodes = ["0", "1", "2"]
    rendezvous = RendezvousHash(nodes=nodes)
    assert 3 == len(rendezvous.nodes)
    assert 1361238019 == rendezvous.hash_function("6666")


@pytest.mark.unit()
def test_seed():
    rendezvous = RendezvousHash(seed=10)
    assert 2981722772 == rendezvous.hash_function("6666")


@pytest.mark.unit()
def test_add_node():
    rendezvous = RendezvousHash()
    rendezvous.add_node("1")

    assert 1 == len(rendezvous.nodes)
    rendezvous.add_node("1")

    assert 1 == len(rendezvous.nodes)
    rendezvous.add_node("2")

    assert 2 == len(rendezvous.nodes)
    rendezvous.add_node("1")

    assert 2 == len(rendezvous.nodes)


@pytest.mark.unit()
def test_remove_node():
    nodes = ["0", "1", "2"]
    rendezvous = RendezvousHash(nodes=nodes)
    rendezvous.remove_node("2")

    assert 2 == len(rendezvous.nodes)

    with pytest.raises(ValueError):
        rendezvous.remove_node("2")

    assert 2 == len(rendezvous.nodes)

    rendezvous.remove_node("1")
    assert 1 == len(rendezvous.nodes)

    rendezvous.remove_node("0")
    assert 0 == len(rendezvous.nodes)


@pytest.mark.unit()
def test_get_node():
    nodes = ["0", "1", "2"]
    rendezvous = RendezvousHash(nodes=nodes)
    assert "0" == rendezvous.get_node("ok")
    assert "1" == rendezvous.get_node("mykey")
    assert "2" == rendezvous.get_node("wat")


@pytest.mark.unit()
def test_get_node_after_removal():
    nodes = ["0", "1", "2"]
    rendezvous = RendezvousHash(nodes=nodes)
    rendezvous.remove_node("1")

    assert "0" == rendezvous.get_node("ok")
    assert "0" == rendezvous.get_node("mykey")
    assert "2" == rendezvous.get_node("wat")


@pytest.mark.unit()
def test_get_node_after_addition():
    nodes = ["0", "1", "2"]
    rendezvous = RendezvousHash(nodes=nodes)
    assert "0" == rendezvous.get_node("ok")
    assert "1" == rendezvous.get_node("mykey")
    assert "2" == rendezvous.get_node("wat")
    assert "2" == rendezvous.get_node("lol")
    rendezvous.add_node("3")

    assert "0" == rendezvous.get_node("ok")
    assert "1" == rendezvous.get_node("mykey")
    assert "2" == rendezvous.get_node("wat")
    assert "3" == rendezvous.get_node("lol")


@pytest.mark.unit()
def test_grow():
    rendezvous = RendezvousHash()

    placements = {}

    for i in range(10):
        rendezvous.add_node(str(i))
        placements[str(i)] = []

    for i in range(1000):
        node = rendezvous.get_node(str(i))
        placements[node].append(i)

    new_placements = {}

    for i in range(20):
        rendezvous.add_node(str(i))
        new_placements[str(i)] = []

    for i in range(1000):
        node = rendezvous.get_node(str(i))
        new_placements[node].append(i)

    keys = [k for sublist in placements.values() for k in sublist]
    new_keys = [k for sublist in new_placements.values() for k in sublist]
    assert sorted(keys) == sorted(new_keys)

    added = 0
    removed = 0

    for node, assignments in new_placements.items():
        after = set(assignments)
        before = set(placements.get(node, []))
        removed += len(before.difference(after))
        added += len(after.difference(before))

    assert added == removed
    assert 1062 == (added + removed)


@pytest.mark.unit()
def test_shrink():
    rendezvous = RendezvousHash()

    placements = {}
    for i in range(10):
        rendezvous.add_node(str(i))
        placements[str(i)] = []

    for i in range(1000):
        node = rendezvous.get_node(str(i))
        placements[node].append(i)

    rendezvous.remove_node("9")
    new_placements = {}
    for i in range(9):
        new_placements[str(i)] = []

    for i in range(1000):
        node = rendezvous.get_node(str(i))
        new_placements[node].append(i)

    keys = [k for sublist in placements.values() for k in sublist]
    new_keys = [k for sublist in new_placements.values() for k in sublist]
    assert sorted(keys) == sorted(new_keys)

    added = 0
    removed = 0
    for node, assignments in placements.items():
        after = set(assignments)
        before = set(new_placements.get(node, []))
        removed += len(before.difference(after))
        added += len(after.difference(before))

    assert added == removed
    assert 202 == (added + removed)


def collide(key, seed):
    return 1337


@pytest.mark.unit()
def test_rendezvous_collision():
    nodes = ["c", "b", "a"]
    rendezvous = RendezvousHash(nodes, hash_function=collide)

    for i in range(1000):
        assert "c" == rendezvous.get_node(i)


@pytest.mark.unit()
def test_rendezvous_names():
    nodes = [1, 2, 3, "a", "b", "lol.wat.com"]
    rendezvous = RendezvousHash(nodes, hash_function=collide)

    for i in range(10):
        assert "lol.wat.com" == rendezvous.get_node(i)

    nodes = [1, "a", "0"]
    rendezvous = RendezvousHash(nodes, hash_function=collide)

    for i in range(10):
        assert "a" == rendezvous.get_node(i)