Spaces:
Running
Running
File size: 2,275 Bytes
b247dc4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
"""Array cache test."""
from pathlib import Path
import numpy as np
import pytest
from manifest.caches.array_cache import ArrayCache
def test_init(tmpdir: Path) -> None:
"""Test cache initialization."""
cache = ArrayCache(Path(tmpdir))
assert (tmpdir / "hash2arrloc.sqlite").exists()
assert cache.cur_file_idx == 0
assert cache.cur_offset == 0
def test_put_get(tmpdir: Path) -> None:
"""Test putting and getting."""
cache = ArrayCache(tmpdir)
cache.max_memmap_size = 5
arr = np.random.rand(10, 10)
with pytest.raises(ValueError) as exc_info:
cache.put("key", arr)
assert str(exc_info.value) == ("Array is too large to be cached. Max is 5")
cache.max_memmap_size = 120
cache.put("key", arr)
assert np.allclose(cache.get("key"), arr)
assert cache.get("key").dtype == arr.dtype
assert cache.cur_file_idx == 0
assert cache.cur_offset == 100
assert cache.hash2arrloc["key"] == {
"file_idx": 0,
"offset": 0,
"flatten_size": 100,
"shape": (10, 10),
"dtype": np.dtype("float64"),
}
arr2 = np.random.randint(0, 3, size=(10, 10))
cache.put("key2", arr2)
assert np.allclose(cache.get("key2"), arr2)
assert cache.get("key2").dtype == arr2.dtype
assert cache.cur_file_idx == 1
assert cache.cur_offset == 100
assert cache.hash2arrloc["key2"] == {
"file_idx": 1,
"offset": 0,
"flatten_size": 100,
"shape": (10, 10),
"dtype": np.dtype("int64"),
}
cache = ArrayCache(tmpdir)
assert cache.hash2arrloc["key"] == {
"file_idx": 0,
"offset": 0,
"flatten_size": 100,
"shape": (10, 10),
"dtype": np.dtype("float64"),
}
assert cache.hash2arrloc["key2"] == {
"file_idx": 1,
"offset": 0,
"flatten_size": 100,
"shape": (10, 10),
"dtype": np.dtype("int64"),
}
assert np.allclose(cache.get("key"), arr)
assert np.allclose(cache.get("key2"), arr2)
def test_contains_key(tmpdir: Path) -> None:
"""Test contains key."""
cache = ArrayCache(tmpdir)
assert not cache.contains_key("key")
arr = np.random.rand(10, 10)
cache.put("key", arr)
assert cache.contains_key("key")
|