mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-11-05 07:39:47 +00:00
297 lines
9.8 KiB
Python
297 lines
9.8 KiB
Python
"""Tests for small utilities in ``reme.utils``."""
|
|
|
|
import asyncio
|
|
import sys
|
|
import threading
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from reme.utils import common_utils
|
|
from reme.utils.counter import (
|
|
COUNTER_LOCK_KEY,
|
|
COUNTER_TREE_KEY,
|
|
global_counter_add,
|
|
global_counter_add_many,
|
|
global_counter_get,
|
|
global_counter_get_all,
|
|
global_counter_inc,
|
|
)
|
|
from reme.utils.similarity_utils import batch_cosine_similarity, cosine_similarity
|
|
|
|
|
|
def test_batch_cosine_similarity_rejects_1d_inputs():
|
|
"""1D vectors should fail with a clear validation error, not IndexError."""
|
|
with pytest.raises(ValueError, match="Expected 2D arrays"):
|
|
batch_cosine_similarity(np.array([1.0, 0.0]), np.array([[1.0, 0.0]]))
|
|
|
|
|
|
def test_batch_cosine_similarity_pairwise_matrix():
|
|
"""Batch cosine returns the full pairwise matrix for valid 2D inputs."""
|
|
result = batch_cosine_similarity(
|
|
np.array([[1.0, 0.0], [0.0, 1.0]]),
|
|
np.array([[1.0, 0.0], [1.0, 1.0]]),
|
|
)
|
|
|
|
assert result.shape == (2, 2)
|
|
np.testing.assert_allclose(result[0], [1.0, 2**-0.5])
|
|
np.testing.assert_allclose(result[1], [0.0, 2**-0.5])
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [np.float16, np.float32, np.float64])
|
|
def test_batch_cosine_similarity_zero_vectors(dtype):
|
|
"""Zero rows and columns have finite zero scores, including stored float16 embeddings."""
|
|
left = np.array([[0, 0], [1, 0], [0, 1]], dtype=dtype)
|
|
right = np.array([[1, 0], [0, 0], [-1, 0]], dtype=dtype)
|
|
|
|
with np.errstate(invalid="raise", divide="raise"):
|
|
result = batch_cosine_similarity(left, right)
|
|
reverse = batch_cosine_similarity(right, left)
|
|
zeros = batch_cosine_similarity(np.zeros((2, 2), dtype=dtype), right)
|
|
all_zeros = batch_cosine_similarity(np.zeros((2, 2), dtype=dtype), np.zeros((3, 2), dtype=dtype))
|
|
|
|
assert result.dtype == dtype
|
|
np.testing.assert_array_equal(result, [[0, 0, 0], [1, 0, -1], [0, 0, 0]])
|
|
np.testing.assert_array_equal(reverse, result.T)
|
|
np.testing.assert_array_equal(zeros, np.zeros((2, 3)))
|
|
np.testing.assert_array_equal(all_zeros, np.zeros((2, 3)))
|
|
|
|
|
|
@pytest.mark.parametrize("left_rows,right_rows", [(0, 2), (2, 0), (0, 0)])
|
|
def test_batch_cosine_similarity_empty_batches(left_rows, right_rows):
|
|
"""Empty batches retain the pairwise output shape."""
|
|
result = batch_cosine_similarity(
|
|
np.zeros((left_rows, 2), dtype=np.float16),
|
|
np.zeros((right_rows, 2), dtype=np.float16),
|
|
)
|
|
assert result.shape == (left_rows, right_rows)
|
|
assert result.dtype == np.float16
|
|
|
|
|
|
def test_batch_cosine_similarity_rejects_mismatched_dimensions():
|
|
"""Batched embeddings must share their feature dimension."""
|
|
with pytest.raises(ValueError, match="Embedding dimensions must match"):
|
|
batch_cosine_similarity(np.zeros((2, 2)), np.zeros((2, 3)))
|
|
|
|
|
|
def test_cosine_similarity_rejects_mismatched_lengths():
|
|
"""Single-vector cosine validates dimensions before computing."""
|
|
with pytest.raises(ValueError, match="Vectors must have same length"):
|
|
cosine_similarity([1.0], [1.0, 2.0])
|
|
|
|
|
|
def test_mock_reme_server_uses_reme_entrypoint(monkeypatch):
|
|
"""The test server helper should spawn the reme CLI module, not legacy reme."""
|
|
captured: dict[str, list[str]] = {}
|
|
|
|
class DummyProcess:
|
|
"""Process stub returned by the patched Popen."""
|
|
|
|
stdout = None
|
|
|
|
def poll(self):
|
|
"""Return a successful process status."""
|
|
return 0
|
|
|
|
def fake_popen(cmd, **_kwargs):
|
|
"""Capture the spawned command."""
|
|
captured["cmd"] = cmd
|
|
return DummyProcess()
|
|
|
|
async def fake_wait_ready(_host, _port, _timeout):
|
|
return None
|
|
|
|
monkeypatch.setattr(common_utils.subprocess, "Popen", fake_popen)
|
|
monkeypatch.setattr(common_utils, "_wait_reme_ready", fake_wait_ready)
|
|
|
|
async def run():
|
|
async with common_utils.mock_reme_server(port=45678, log_to_file=False, enable_logo=False):
|
|
pass
|
|
|
|
asyncio.run(run())
|
|
|
|
assert captured["cmd"][:4] == [sys.executable, "-m", "reme.reme", "start"]
|
|
|
|
|
|
def test_inc_returns_old_value_starting_at_zero():
|
|
"""First call returns 0, then values increase by 1 per call."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_inc(metadata, ["a"]) == 0
|
|
assert global_counter_inc(metadata, ["a"]) == 1
|
|
assert global_counter_inc(metadata, ["a"]) == 2
|
|
|
|
|
|
def test_add_returns_old_value_and_adds_val():
|
|
"""``add`` is fetch-and-add: old value out, ``val`` added in."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_add(metadata, ["a"], 10) == 0
|
|
assert global_counter_add(metadata, ["a"], 5) == 10
|
|
assert global_counter_inc(metadata, ["a"]) == 15
|
|
assert global_counter_get(metadata, ["a"]) == 16
|
|
|
|
|
|
def test_add_many_returns_old_values_and_updates_all_paths():
|
|
"""``add_many`` updates sibling metrics under one counter-tree lock."""
|
|
metadata: dict = {}
|
|
global_counter_add(metadata, ["usage", "input"], 2)
|
|
|
|
previous = global_counter_add_many(
|
|
metadata,
|
|
{
|
|
("usage", "input"): 3,
|
|
("usage", "output"): 4,
|
|
("usage", "total"): 7,
|
|
},
|
|
)
|
|
|
|
assert previous == {
|
|
("usage", "input"): 2,
|
|
("usage", "output"): 0,
|
|
("usage", "total"): 0,
|
|
}
|
|
assert global_counter_get(metadata, ["usage", "input"]) == 5
|
|
assert global_counter_get(metadata, ["usage", "output"]) == 4
|
|
assert global_counter_get(metadata, ["usage", "total"]) == 7
|
|
|
|
|
|
def test_add_many_holds_the_counter_lock_once_for_the_batch():
|
|
"""A multi-metric update is one critical section, not several writes."""
|
|
|
|
class CountingLock:
|
|
"""Lock test double that counts entered critical sections."""
|
|
|
|
def __init__(self):
|
|
self.lock = threading.Lock()
|
|
self.entries = 0
|
|
|
|
def __enter__(self):
|
|
self.lock.acquire()
|
|
self.entries += 1
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
self.lock.release()
|
|
|
|
lock = CountingLock()
|
|
metadata = {COUNTER_LOCK_KEY: lock}
|
|
|
|
global_counter_add_many(
|
|
metadata,
|
|
{
|
|
("usage", "input"): 3,
|
|
("usage", "output"): 4,
|
|
("usage", "total"): 7,
|
|
},
|
|
)
|
|
|
|
assert lock.entries == 1
|
|
|
|
|
|
def test_add_many_validates_every_update_before_mutating():
|
|
"""One invalid update leaves every valid counter unchanged."""
|
|
metadata: dict = {}
|
|
|
|
with pytest.raises(TypeError, match="increments must be integers"):
|
|
global_counter_add_many(
|
|
metadata,
|
|
{
|
|
("usage", "input"): 3,
|
|
("usage", "output"): "invalid",
|
|
},
|
|
)
|
|
|
|
assert COUNTER_TREE_KEY not in metadata
|
|
|
|
|
|
def test_counters_are_isolated_by_key_path():
|
|
"""Sibling and nested keys, plus the root, hold independent counters."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_inc(metadata, ["a"]) == 0
|
|
assert global_counter_inc(metadata, ["b"]) == 0
|
|
assert global_counter_inc(metadata, ["a", "child"]) == 0
|
|
assert global_counter_inc(metadata, []) == 0
|
|
|
|
assert global_counter_get(metadata, ["a"]) == 1
|
|
assert global_counter_get(metadata, ["b"]) == 1
|
|
assert global_counter_get(metadata, ["a", "child"]) == 1
|
|
assert global_counter_get(metadata, []) == 1
|
|
|
|
|
|
def test_get_does_not_create_missing_nodes():
|
|
"""``get`` reports 0 for missing paths and leaves the tree untouched."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_get(metadata, ["missing"]) == 0
|
|
assert COUNTER_TREE_KEY not in metadata
|
|
|
|
global_counter_inc(metadata, ["a"])
|
|
assert global_counter_get(metadata, ["a", "missing"]) == 0
|
|
assert "missing" not in metadata[COUNTER_TREE_KEY]["children"]["a"]["children"]
|
|
|
|
|
|
def test_get_all_returns_none_for_missing_key():
|
|
"""``get_all`` returns None when the tree or the path does not exist."""
|
|
metadata: dict = {}
|
|
|
|
assert global_counter_get_all(metadata, []) is None
|
|
assert global_counter_get_all(metadata, ["missing"]) is None
|
|
|
|
global_counter_inc(metadata, ["a"])
|
|
assert global_counter_get_all(metadata, ["missing"]) is None
|
|
assert global_counter_get_all(metadata, ["a", "missing"]) is None
|
|
|
|
|
|
def test_get_all_returns_subtree_and_whole_tree():
|
|
"""``get_all`` returns the node at ``key``; an empty key returns the root."""
|
|
metadata: dict = {}
|
|
global_counter_add(metadata, ["a"], 2)
|
|
global_counter_add(metadata, ["a", "child"], 3)
|
|
|
|
subtree = global_counter_get_all(metadata, ["a"])
|
|
assert subtree == {"value": 2, "children": {"child": {"value": 3, "children": {}}}}
|
|
|
|
root = global_counter_get_all(metadata, [])
|
|
assert root["value"] == 0
|
|
assert root["children"]["a"] == subtree
|
|
|
|
|
|
def test_get_all_returns_deep_copy():
|
|
"""Mutating the returned subtree must not affect the live counter tree."""
|
|
metadata: dict = {}
|
|
global_counter_add(metadata, ["a", "child"], 3)
|
|
|
|
subtree = global_counter_get_all(metadata, ["a"])
|
|
subtree["value"] = 999
|
|
subtree["children"]["child"]["value"] = 999
|
|
subtree["children"]["extra"] = {"value": 1, "children": {}}
|
|
|
|
assert global_counter_get(metadata, ["a"]) == 0
|
|
assert global_counter_get(metadata, ["a", "child"]) == 3
|
|
assert global_counter_get_all(metadata, ["a", "extra"]) is None
|
|
|
|
|
|
def test_concurrent_inc_yields_unique_values():
|
|
"""Parallel increments on one key never return duplicate values."""
|
|
metadata: dict = {}
|
|
results: list[int] = []
|
|
results_lock = threading.Lock()
|
|
calls_per_thread = 200
|
|
|
|
def worker():
|
|
for _ in range(calls_per_thread):
|
|
value = global_counter_inc(metadata, ["shared"])
|
|
with results_lock:
|
|
results.append(value)
|
|
|
|
threads = [threading.Thread(target=worker) for _ in range(8)]
|
|
for thread in threads:
|
|
thread.start()
|
|
for thread in threads:
|
|
thread.join()
|
|
|
|
total = len(threads) * calls_per_thread
|
|
assert sorted(results) == list(range(total))
|
|
assert global_counter_get(metadata, ["shared"]) == total
|