hivemind/tests/test_dht_storage.py

import time

from hivemind.utils.timed_storage import get_dht_time
from hivemind.dht.storage import DHTLocalStorage, DHTID, DictionaryDHTValue
from hivemind.utils.serializer import MSGPackSerializer


def test_store():
    d = DHTLocalStorage()
    d.store(DHTID.generate("key"), b"val", get_dht_time() + 0.5)
    assert d.get(DHTID.generate("key"))[0] == b"val", "Wrong value"


def test_get_expired():
    d = DHTLocalStorage()
    d.store(DHTID.generate("key"), b"val", get_dht_time() + 0.1)
    time.sleep(0.5)
    assert d.get(DHTID.generate("key")) is None, "Expired value must be deleted"


def test_get_empty():
    d = DHTLocalStorage()
    assert d.get(DHTID.generate(source="key")) is None, "DHTLocalStorage returned non-existent value"


def test_change_expiration_time():
    d = DHTLocalStorage()
    d.store(DHTID.generate("key"), b"val1", get_dht_time() + 1)
    assert d.get(DHTID.generate("key"))[0] == b"val1", "Wrong value"
    d.store(DHTID.generate("key"), b"val2", get_dht_time() + 200)
    time.sleep(1)
    assert d.get(DHTID.generate("key"))[0] == b"val2", "Value must be changed, but still kept in table"


def test_maxsize_cache():
    d = DHTLocalStorage(maxsize=2)
    d.store(DHTID.generate("key1a"), b"val1a", get_dht_time() + 1)
    d.store(DHTID.generate("key1b"), b"val1b", get_dht_time() + 1)
    d.store(DHTID.generate("key1a"), b"val1a2", get_dht_time() + 2)
    d.store(DHTID.generate("key2"), b"val2", get_dht_time() + 200)
    assert d.get(DHTID.generate("key2"))[0] == b"val2", "Value with bigger exp. time must be kept"
    assert d.get(DHTID.generate("key1a"))[0] == b"val1a2", "Value with bigger exp. time must be kept"
    assert d.get(DHTID.generate("key1b")) is None, "Value with less exp time, must be deleted"


def test_localstorage_top():
    d = DHTLocalStorage(maxsize=3)
    d.store(DHTID.generate("key1"), b"val1", get_dht_time() + 1)
    d.store(DHTID.generate("key2"), b"val2", get_dht_time() + 2)
    d.store(DHTID.generate("key3"), b"val3", get_dht_time() + 4)
    assert d.top()[0] == DHTID.generate("key1") and d.top()[1].value == b"val1"

    d.store(DHTID.generate("key1"), b"val1_new", get_dht_time() + 3)
    assert d.top()[0] == DHTID.generate("key2") and d.top()[1].value == b"val2"

    del d[DHTID.generate("key2")]
    assert d.top()[0] == DHTID.generate("key1") and d.top()[1].value == b"val1_new"
    d.store(DHTID.generate("key2"), b"val2_new", get_dht_time() + 5)
    d.store(DHTID.generate("key4"), b"val4", get_dht_time() + 6)  # key4 will push out key1 due to maxsize

    assert d.top()[0] == DHTID.generate("key3") and d.top()[1].value == b"val3"


def test_localstorage_nested():
    time = get_dht_time()
    d1 = DHTLocalStorage()
    d2 = DictionaryDHTValue()
    d2.store("subkey1", b"value1", time + 2)
    d2.store("subkey2", b"value2", time + 3)
    d2.store("subkey3", b"value3", time + 1)

    assert d2.latest_expiration_time == time + 3
    for subkey, (subvalue, subexpiration) in d2.items():
        assert d1.store_subkey(DHTID.generate("foo"), subkey, subvalue, subexpiration)
    assert d1.store(DHTID.generate("bar"), b"456", time + 2)
    assert d1.get(DHTID.generate("foo"))[0].data == d2.data
    assert d1.get(DHTID.generate("foo"))[1] == d2.latest_expiration_time
    assert d1.get(DHTID.generate("foo"))[0].get("subkey1") == (b"value1", time + 2)
    assert len(d1.get(DHTID.generate("foo"))[0]) == 3
    assert d1.store_subkey(DHTID.generate("foo"), "subkey4", b"value4", time + 4)
    assert len(d1.get(DHTID.generate("foo"))[0]) == 4

    assert (
        d1.store_subkey(DHTID.generate("bar"), "subkeyA", b"valueA", time + 1) is False
    )  # prev has better expiration
    assert d1.store_subkey(DHTID.generate("bar"), "subkeyA", b"valueA", time + 3)  # new value has better expiration
    assert d1.store_subkey(DHTID.generate("bar"), "subkeyB", b"valueB", time + 4)  # new value has better expiration
    assert d1.store_subkey(DHTID.generate("bar"), "subkeyA", b"valueA+", time + 5)  # overwrite subkeyA under key bar
    assert all(subkey in d1.get(DHTID.generate("bar"))[0] for subkey in ("subkeyA", "subkeyB"))
    assert len(d1.get(DHTID.generate("bar"))[0]) == 2 and d1.get(DHTID.generate("bar"))[1] == time + 5

    assert d1.store(DHTID.generate("foo"), b"nothing", time + 3.5) is False  # previous value has better expiration
    assert d1.get(DHTID.generate("foo"))[0].get("subkey2") == (b"value2", time + 3)
    assert d1.store(DHTID.generate("foo"), b"nothing", time + 5) is True  # new value has better expiraiton
    assert d1.get(DHTID.generate("foo")) == (b"nothing", time + 5)  # value should be replaced


def test_localstorage_freeze():
    d = DHTLocalStorage(maxsize=2)

    with d.freeze():
        d.store(DHTID.generate("key1"), b"val1", get_dht_time() + 0.01)
        assert DHTID.generate("key1") in d
        time.sleep(0.03)
        assert DHTID.generate("key1") in d
    assert DHTID.generate("key1") not in d

    with d.freeze():
        d.store(DHTID.generate("key1"), b"val1", get_dht_time() + 1)
        d.store(DHTID.generate("key2"), b"val2", get_dht_time() + 2)
        d.store(DHTID.generate("key3"), b"val3", get_dht_time() + 3)  # key3 will push key1 out due to maxsize
        assert DHTID.generate("key1") in d
    assert DHTID.generate("key1") not in d


def test_localstorage_serialize():
    d1 = DictionaryDHTValue()
    d2 = DictionaryDHTValue()

    now = get_dht_time()
    d1.store("key1", b"ololo", now + 1)
    d2.store("key2", b"pysh", now + 1)
    d2.store("key3", b"pyshpysh", now + 2)

    data = MSGPackSerializer.dumps([d1, d2, 123321])
    assert isinstance(data, bytes)
    new_d1, new_d2, new_value = MSGPackSerializer.loads(data)
    assert isinstance(new_d1, DictionaryDHTValue) and isinstance(new_d2, DictionaryDHTValue) and new_value == 123321
    assert "key1" in new_d1 and len(new_d1) == 1
    assert "key1" not in new_d2 and len(new_d2) == 2 and new_d2.get("key3") == (b"pyshpysh", now + 2)