# map/tests/test_servecache.py import os import tempfile import time import unittest from pathlib import Path from unittest import mock import numpy as np import servecache as SC class ServeCacheTest(unittest.TestCase): def setUp(self): self.tmp = Path(tempfile.mkdtemp()) self.arrays = {"a": np.arange(10, dtype=np.float32), "ids": np.array([3, 1, 2], np.uint64), "m": np.array([True, False]), "xyz": np.random.default_rng(1).normal(size=(5, 3))} def test_round_trip_keeps_values_and_dtypes(self): d = SC.cache_dir(self.tmp, "abc") self.assertEqual(d.name, f"serve-abc-v{SC.VERSION}") SC.write(d, self.arrays, {"note": "x"}) arrays, meta = SC.read(d) self.assertEqual(meta["note"], "x") for k, v in self.arrays.items(): self.assertIsInstance(arrays[k], np.memmap) self.assertEqual(arrays[k].dtype, v.dtype) np.testing.assert_array_equal(arrays[k], v) def test_read_loads_more_arrays_than_the_soft_file_limit(self): import resource soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) if hard != resource.RLIM_INFINITY and hard < 400: self.skipTest("hard open-file limit too low") d = SC.cache_dir(self.tmp, "many") SC.write(d, {f"f{i}": np.arange(3) + i for i in range(300)}, {}) got = None try: resource.setrlimit(resource.RLIMIT_NOFILE, (128, hard)) # each memory map keeps a file open got = SC.read(d) self.assertIsNotNone(got, "a complete cache must not read as damaged") self.assertEqual(int(got[0]["f299"][2]), 301) finally: del got resource.setrlimit(resource.RLIMIT_NOFILE, (soft, hard)) def test_read_ignores_a_folder_without_done(self): d = SC.cache_dir(self.tmp, "abc") SC.write(d, self.arrays, {}) (d / "done").unlink() self.assertIsNone(SC.read(d)) self.assertIsNone(SC.read(self.tmp / "missing")) def test_failed_write_leaves_no_final_folder(self): d = SC.cache_dir(self.tmp, "abc") with mock.patch.object(SC, "save_npy", side_effect=OSError("disk full")): with self.assertRaises(OSError): SC.write(d, self.arrays, {}) self.assertFalse(d.exists()) self.assertEqual(list(self.tmp.iterdir()), [], "the temporary folder is removed too") def test_load_or_build_builds_once(self): d = SC.cache_dir(self.tmp, "abc") calls = [] def build(): calls.append(1) return self.arrays, {"n": 1} a1, _ = SC.load_or_build(d, build) a2, meta = SC.load_or_build(d, build) self.assertEqual(len(calls), 1) self.assertEqual(meta["n"], 1) self.assertIsInstance(a1["a"], np.memmap, "even the first load hands back the written maps") np.testing.assert_array_equal(a2["xyz"], self.arrays["xyz"]) def test_unwritable_falls_back_to_ram(self): d = SC.cache_dir(self.tmp, "abc") logs = [] with mock.patch.object(SC, "write", side_effect=OSError("read-only file system")): arrays, _ = SC.load_or_build(d, lambda: (self.arrays, {}), log=logs.append) self.assertIs(arrays["a"], self.arrays["a"]) self.assertEqual(len(logs), 1) self.assertIn("read-only file system", logs[0]) def test_prune_removes_only_unkept_serve_folders(self): keep = SC.cache_dir(self.tmp, "new") old = SC.cache_dir(self.tmp, "old") for d in (keep, old): SC.write(d, self.arrays, {}) area = self.tmp / "0123abcd" area.mkdir() (area / "cells.npz").write_bytes(b"x") tiles = self.tmp / "v4-abc" tiles.mkdir() gone = SC.prune(self.tmp, [keep]) self.assertEqual(gone, [old]) self.assertTrue(keep.exists() and area.exists() and tiles.exists()) def test_prune_keeps_young_temporary_folders_and_removes_old_ones(self): young = self.tmp / f"serve-x-v{SC.VERSION}.tmp-1" stale = self.tmp / f"serve-y-v{SC.VERSION}.tmp-2" young.mkdir() stale.mkdir() t = time.time() - 7200 os.utime(stale, (t, t)) SC.prune(self.tmp, []) self.assertTrue(young.exists(), "may belong to a server still writing it") self.assertFalse(stale.exists()) def test_prune_removes_a_young_temporary_folder_of_a_process_that_is_gone(self): import subprocess pr = subprocess.Popen(["true"]) pr.wait() orphan = self.tmp / f"serve-z-v{SC.VERSION}.tmp-{pr.pid}" orphan.mkdir() SC.prune(self.tmp, []) self.assertFalse(orphan.exists(), "its writer died: nobody finishes it") def test_prune_skips_a_folder_another_process_removed_meanwhile(self): gone = self.tmp / f"serve-q-v{SC.VERSION}.tmp-1" with mock.patch.object(Path, "iterdir", lambda _: iter([gone])), \ mock.patch.object(Path, "is_dir", lambda _: True): self.assertEqual(SC.prune(self.tmp, []), []) def test_another_server_finishing_first_is_no_failure(self): import errno d = SC.cache_dir(self.tmp, "race") real = os.rename def other_first(src, dst): with mock.patch.object(SC.os, "rename", real): SC.write(Path(dst), self.arrays, {}) raise OSError(errno.ENOTEMPTY, "Directory not empty") logs = [] with mock.patch.object(SC.os, "rename", other_first): arrays, _ = SC.load_or_build(d, lambda: (self.arrays, {}), log=logs.append) self.assertEqual(logs, [], "no RAM fallback") self.assertIsInstance(arrays["a"], np.memmap) self.assertFalse([p for p in self.tmp.iterdir() if ".tmp-" in p.name]) def test_a_damaged_cache_is_rebuilt(self): for damage in ("empty", "truncated"): d = SC.cache_dir(self.tmp, damage) SC.write(d, self.arrays, {}) f = d / "xyz.npy" f.write_bytes(b"" if damage == "empty" else f.read_bytes()[:200]) # e.g. power lost after the rename logs = [] arrays, _ = SC.load_or_build(d, lambda: (self.arrays, {}), log=logs.append) self.assertIsInstance(arrays["xyz"], np.memmap, damage) np.testing.assert_array_equal(arrays["xyz"], self.arrays["xyz"]) self.assertIsNotNone(SC.read(d), f"{damage}: healed for the next start") if __name__ == "__main__": unittest.main() class SaveNpyTest(unittest.TestCase): def test_same_bytes_as_np_save(self): rng = np.random.default_rng(0) arrays = [rng.normal(size=3_000_000), rng.random((1000, 3)), np.zeros(0), np.array(3.5), rng.integers(0, 9, 50_000).astype(np.int8), rng.random(10_000) < 0.5, np.arange(12, dtype=np.uint64).reshape(3, 4)[:, ::2], np.asfortranarray(rng.random((300, 7)))] with tempfile.TemporaryDirectory() as t: for i, a in enumerate(arrays): with self.subTest(i=i, dtype=str(a.dtype), shape=a.shape): np.save(Path(t) / "want.npy", np.ascontiguousarray(a), allow_pickle=False) SC.save_npy(Path(t) / "got.npy", a) self.assertEqual((Path(t) / "got.npy").read_bytes(), (Path(t) / "want.npy").read_bytes()) def test_written_in_plain_writes_not_tofile(self): """numpy's tofile() path preallocates the file; btrfs then stores it uncompressed.""" from numpy.lib import format as F seen = [] real = F.write_array def spy(fp, *a, **k): seen.append(F.isfileobj(fp)) return real(fp, *a, **k) with tempfile.TemporaryDirectory() as t, mock.patch.object(F, "write_array", spy): SC.save_npy(Path(t) / "a.npy", np.ones(1000)) self.assertTrue(np.array_equal(np.load(Path(t) / "a.npy"), np.ones(1000))) self.assertEqual(seen, [False])