Spaces:
Running
Running
Download tests/test_memory.py from edbeeching/agent-memory: direct link, hf CLI and curl.
- Browser
- Download file 17.2 kB
-
https://huggingface.co/spaces/edbeeching/agent-memory/resolve/main/tests/test_memory.py
- Command line
-
hf download hf://spaces/edbeeching/agent-memory/tests/test_memory.py
-
curl -L -o test_memory.py https://huggingface.co/spaces/edbeeching/agent-memory/resolve/main/tests/test_memory.py
17.2 kB
| import importlib.util | |
| import json | |
| from pathlib import Path | |
| import subprocess | |
| import sys | |
| from concurrent.futures import ThreadPoolExecutor | |
| from types import SimpleNamespace | |
| import pytest | |
| ROOT = Path(__file__).resolve().parents[1] | |
| spec = importlib.util.spec_from_file_location("memory", ROOT / "memory-use/scripts/memory.py") | |
| m = importlib.util.module_from_spec(spec) | |
| spec.loader.exec_module(m) | |
| class Bucket: | |
| """An in-memory fake of bucket operations; never used by the production CLI.""" | |
| def __init__(self): | |
| self.buckets = {} | |
| self.calls = [] | |
| self.fail = None | |
| self.fail_marker = False | |
| self.unknown_put = False | |
| def call(self, op, **kw): | |
| self.calls.append((op, kw)) | |
| if self.fail: | |
| raise m.MemoryError("Simulated " + self.fail, self.fail) | |
| if op == "whoami": | |
| return {"name": "owner", "orgs": ["team"]} | |
| if op == "create": | |
| if kw["bucket"] in self.buckets: | |
| raise m.MemoryError("Already exists", "conflict") | |
| self.buckets[kw["bucket"]] = {"private": True, "files": {}} | |
| return {} | |
| if op == "core": | |
| return {"path": str(ROOT / "memory-use")} | |
| if kw["bucket"] not in self.buckets: | |
| raise m.MemoryError("Missing", "not-found") | |
| bucket = self.buckets[kw["bucket"]] | |
| files = bucket["files"] | |
| if op == "info": | |
| return {"id": kw["bucket"], "private": bucket["private"]} | |
| if op == "put": | |
| if self.fail_marker and kw["path"].endswith("/SKILL.md"): | |
| raise m.MemoryError("Interrupted publication", "offline") | |
| files[kw["path"]] = kw["text"] | |
| if self.unknown_put: | |
| raise m.MemoryError("Lost acknowledgment", "offline") | |
| return {} | |
| if op == "get": | |
| if kw["path"] not in files: | |
| raise m.MemoryError("Missing", "not-found") | |
| return {"text": files[kw["path"]]} | |
| if op == "list": | |
| return [{"path": p, "size": len(t.encode())} for p, t in files.items() if p.startswith(kw["prefix"])] | |
| raise AssertionError(op) | |
| def setup(tmp_path): | |
| store = Bucket() | |
| client = m.Memory(tmp_path / "state", store) | |
| work = tmp_path / "project" | |
| work.mkdir() | |
| client.provision("project", work) | |
| return client, store, work | |
| def example_skill(tmp_path, body="Use a release checklist."): | |
| folder = tmp_path / "release-check" | |
| folder.mkdir(exist_ok=True) | |
| (folder / "SKILL.md").write_text("---\nname: release-check\ndescription: Check releases.\n---\n\n" + body) | |
| (folder / "references").mkdir(exist_ok=True) | |
| (folder / "references/check.md").write_text("Verify the release artifact.") | |
| return folder | |
| def setup_runtime(client, tmp_path, store, work): | |
| instructions = tmp_path / "AGENTS.md" | |
| instructions.write_text("Keep existing conventions.\n") | |
| return client.bootstrap("project", work, bucket="owner/agent-memory", skills_dir=tmp_path / "skills", instructions=instructions) | |
| def test_two_machines_share_memories_and_keep_corrections(setup, tmp_path): | |
| one, store, work = setup | |
| saved = one.save("Use uv for this project.", "Package tooling", cwd=work) | |
| two = m.Memory(tmp_path / "second-machine", store) | |
| other_checkout = tmp_path / "other-checkout" | |
| other_checkout.mkdir() | |
| two.provision("project", other_checkout, "owner/agent-memory") | |
| results = two.search("UV", cwd=other_checkout) | |
| assert results["sync"] == "fresh" and len(results["results"]) == 1 | |
| assert "Use uv" in two.read(saved["id"], cwd=other_checkout)["text"] | |
| correction = two.save("Correction: use pip. Supersedes " + saved["id"], "Tooling correction", cwd=other_checkout) | |
| assert saved["id"] != correction["id"] | |
| assert len(one.search(cwd=work)["results"]) == 2 | |
| def test_unknown_upload_retries_identical_bytes_after_configuration_change(setup): | |
| client, store, work = setup | |
| store.unknown_put = True | |
| result = client.save("Verified result", "Checkpoint", kind="checkpoint", cwd=work) | |
| assert result["status"] == "pending" | |
| original = dict(store.buckets["owner/agent-memory"]["files"]) | |
| store.unknown_put = False | |
| client.provision("other", work, create_bucket="owner/other") | |
| assert client.retry(result["id"])["status"] == "saved" | |
| assert store.buckets["owner/agent-memory"]["files"] == original | |
| assert not store.buckets["owner/other"]["files"] | |
| assert not client.pending() | |
| def test_save_offline_is_durable_and_retryable(setup): | |
| client, store, work = setup | |
| store.fail = "offline" | |
| result = client.save("Keep this finding", "Finding", cwd=work) | |
| assert result["status"] == "pending" | |
| assert client.pending()[0]["text"].endswith("Keep this finding\n") | |
| restarted = m.Memory(client.home, store) | |
| store.fail = None | |
| assert restarted.retry(result["id"])["status"] == "saved" | |
| def test_no_cleanup_before_remote_confirmation(setup, monkeypatch): | |
| client, store, work = setup | |
| original = m.write_json | |
| def crash_on_receipt(path, value): | |
| if Path(path).name == "last-upload.json": | |
| raise OSError("Simulated disk failure after upload") | |
| return original(path, value) | |
| monkeypatch.setattr(m, "write_json", crash_on_receipt) | |
| with pytest.raises(OSError): | |
| client.save("Retain until acknowledged", "Crash case", cwd=work) | |
| job = client.pending()[0] | |
| assert len(store.buckets["owner/agent-memory"]["files"]) == 1 | |
| monkeypatch.setattr(m, "write_json", original) | |
| client.retry(job["id"]) | |
| assert len(store.buckets["owner/agent-memory"]["files"]) == 1 | |
| assert not client.pending() | |
| def test_cached_search_distinguishes_failed_sync_and_does_not_prune(setup, failure): | |
| client, store, work = setup | |
| client.save("Cached evidence", "Evidence", cwd=work) | |
| client.search(cwd=work) | |
| store.fail = failure | |
| result = client.search("evidence", cwd=work) | |
| assert result["stale"] and len(result["results"]) == 1 | |
| assert result["sync"] == ("auth-error" if failure == "auth" else "offline") | |
| def test_incomplete_listing_and_download_preserve_cache(setup, monkeypatch): | |
| client, store, work = setup | |
| old = client.save("Keep cached note", "Old", cwd=work) | |
| client.search(cwd=work) | |
| original = store.call | |
| def interrupted(op, **kw): | |
| if op == "list": | |
| raise m.MemoryError("Pagination failed after page one", "offline") | |
| return original(op, **kw) | |
| monkeypatch.setattr(store, "call", interrupted) | |
| assert len(client.search(cwd=work)["results"]) == 1 | |
| monkeypatch.setattr(store, "call", original) | |
| store.buckets["owner/agent-memory"]["files"].pop(old["path"]) | |
| new = client.save("New note", "New", cwd=work) | |
| def broken_download(op, **kw): | |
| if op == "get" and kw["path"] == new["path"]: | |
| raise m.MemoryError("Download interrupted", "offline") | |
| return original(op, **kw) | |
| monkeypatch.setattr(store, "call", broken_download) | |
| assert client.search(cwd=work)["stale"] | |
| assert (client.cache("owner/agent-memory") / old["path"]).exists() | |
| assert not (client.cache("owner/agent-memory") / new["path"]).exists() | |
| monkeypatch.setattr(store, "call", original) | |
| assert not client.search(cwd=work)["stale"] | |
| assert not (client.cache("owner/agent-memory") / old["path"]).exists() | |
| def test_explicit_missing_bucket_never_creates_replacement(tmp_path): | |
| store = Bucket() | |
| client = m.Memory(tmp_path / "state", store) | |
| with pytest.raises(m.MemoryError): | |
| client.provision("project", tmp_path, "owner/missing") | |
| assert not any(op == "create" for op, _ in store.calls) | |
| def test_configured_missing_bucket_is_not_replaced(setup): | |
| client, store, work = setup | |
| store.buckets.clear() | |
| store.calls.clear() | |
| with pytest.raises(m.MemoryError): | |
| client.provision("project", work) | |
| assert not any(op == "create" for op, _ in store.calls) | |
| def test_public_default_bucket_is_not_reconfigured(tmp_path): | |
| store = Bucket() | |
| store.buckets["owner/agent-memory"] = {"private": False, "files": {}} | |
| with pytest.raises(m.MemoryError, match="private"): | |
| m.Memory(tmp_path / "state", store).provision("project", tmp_path) | |
| assert store.buckets["owner/agent-memory"]["private"] is False | |
| assert not any(op == "create" for op, _ in store.calls) | |
| def test_longest_project_ancestor_and_directory_boundaries(setup, tmp_path): | |
| client, _, work = setup | |
| nested = work / "nested" | |
| nested.mkdir() | |
| client.provision("nested", nested, "owner/agent-memory") | |
| assert client.project(nested / "src")["project"] == "nested" | |
| assert client.project(work / "other")["project"] == "project" | |
| with pytest.raises(m.MemoryError, match="no memory project"): | |
| client.save("Do not route incorrectly", "Unmapped", cwd=tmp_path / "project-other") | |
| def test_bootstrap_repeat_preserves_instructions_and_verification(setup, tmp_path): | |
| client, store, work = setup | |
| first = setup_runtime(client, tmp_path, store, work) | |
| instructions = tmp_path / "AGENTS.md" | |
| instructions.write_text(instructions.read_text() + "\nMore unrelated instructions.\n") | |
| second = client.bootstrap("project", work, bucket="owner/agent-memory", skills_dir=tmp_path / "skills", instructions=instructions) | |
| assert first["verification_id"] == second["verification_id"] | |
| assert len(store.buckets["owner/agent-memory"]["files"]) == 1 | |
| text = instructions.read_text() | |
| assert text.startswith("Keep existing conventions.") and text.endswith("More unrelated instructions.\n") | |
| assert text.count(m.BEGIN) == 1 | |
| def test_skill_revisions_partial_publication_and_frozen_retry(setup, tmp_path): | |
| client, store, work = setup | |
| folder = example_skill(tmp_path) | |
| store.fail_marker = True | |
| first = client.publish(folder, cwd=work) | |
| assert first["status"] == "pending" | |
| assert not client.search(skills=True, cwd=work)["results"] | |
| original_text = (folder / "SKILL.md").read_text() | |
| (folder / "SKILL.md").write_text(original_text + "\nA subsequent local edit.") | |
| store.fail_marker = False | |
| client.retry(first["id"]) | |
| assert store.buckets["owner/agent-memory"]["files"][first["path"] + "SKILL.md"] == original_text | |
| second = client.publish(folder, cwd=work) | |
| assert second["revision"] != first["revision"] | |
| assert len(client.search(skills=True, cwd=work)["results"]) == 2 | |
| def test_install_preserves_modified_skill_and_explicit_replacement_backs_up(setup, tmp_path): | |
| client, store, work = setup | |
| setup_runtime(client, tmp_path, store, work) | |
| folder = example_skill(tmp_path) | |
| published = client.publish(folder, cwd=work) | |
| client.install_shared("release-check", published["revision"], cwd=work) | |
| installed = tmp_path / "skills/release-check/references/check.md" | |
| installed.write_text("Local modifications") | |
| with pytest.raises(m.MemoryError, match="Existing skill differs"): | |
| client.install_shared("release-check", published["revision"], cwd=work) | |
| assert installed.read_text() == "Local modifications" | |
| result = client.install_shared("release-check", published["revision"], cwd=work, replace=True) | |
| assert (Path(result["backup"]) / "references/check.md").read_text() == "Local modifications" | |
| assert "release-check" in (tmp_path / "AGENTS.md").read_text() | |
| def test_unsafe_paths_rejected(bad): | |
| with pytest.raises(m.MemoryError): | |
| m.validate_skill({"SKILL.md": "---\nname: example\ndescription: Example\n---\nBody", bad: "bad"}) | |
| def test_symlink_publish_and_install_destination_rejected(setup, tmp_path): | |
| client, store, work = setup | |
| folder = example_skill(tmp_path) | |
| (folder / "link").symlink_to(tmp_path / "outside") | |
| with pytest.raises(m.MemoryError, match="symlink"): | |
| client.publish(folder, cwd=work) | |
| (folder / "link").unlink() | |
| root = tmp_path / "linked-skills" | |
| root.symlink_to(folder, target_is_directory=True) | |
| with pytest.raises(m.MemoryError, match="Symlinks"): | |
| client.install_files(m.skill_files(folder), "release-check", root, {}) | |
| def test_core_name_cannot_be_published_or_installed_from_bucket(setup): | |
| client, _, work = setup | |
| with pytest.raises(m.MemoryError, match="pinned Space"): | |
| client.publish(ROOT / "memory-use", cwd=work) | |
| with pytest.raises(m.MemoryError, match="cannot replace"): | |
| client.install_shared("memory-use", "abc", cwd=work) | |
| def test_reminder_has_no_network_or_transcript_reads(setup, tmp_path): | |
| client, store, work = setup | |
| store.calls.clear() | |
| payload = {"hook_event_name": "UserPromptSubmit", "cwd": str(work), "transcript_path": "/must/not/read", "prompt": "Do not echo my message"} | |
| result = m.hook_output(client, payload) | |
| assert result["hookSpecificOutput"]["hookEventName"] == "UserPromptSubmit" | |
| assert "Do not echo" not in json.dumps(result) | |
| assert not store.calls | |
| assert m.hook_output(client, {**payload, "cwd": str(tmp_path / "unmapped")}) == {} | |
| def test_hook_merge_remove_preserves_unrelated_handlers(setup, tmp_path, monkeypatch): | |
| client, store, work = setup | |
| setup_runtime(client, tmp_path, store, work) | |
| config = tmp_path / "hooks.json" | |
| unrelated = {"type": "command", "command": "echo unrelated"} | |
| initial = {"description": "Existing hooks", "hooks": {"SessionStart": [{"matcher": "resume", "hooks": [unrelated]}], "Stop": [{"hooks": [unrelated]}]}} | |
| config.write_text(json.dumps(initial)) | |
| monkeypatch.setattr(m.subprocess, "run", lambda *a, **k: SimpleNamespace(stdout="codex-cli 0.162.1")) | |
| m.hooks(client, "install", config) | |
| m.hooks(client, "install", config) | |
| value = json.loads(config.read_text()) | |
| assert len(value["hooks"]["SessionStart"]) == 2 | |
| m.hooks(client, "remove", config) | |
| assert json.loads(config.read_text()) == initial | |
| def test_credentials_rejected_and_errors_do_not_echo_sdk_secrets(setup, monkeypatch): | |
| client, _, work = setup | |
| token = "hf_" + "x" * 30 | |
| monkeypatch.setenv("HF_TOKEN", token) | |
| with pytest.raises(m.MemoryError, match="credential"): | |
| client.save(token, "Bad note", cwd=work) | |
| exc = Exception("Authorization: " + token) | |
| assert token not in str(m.error_from_exception(exc)) | |
| def test_cli_json_errors_and_hook_fail_open(tmp_path): | |
| script = ROOT / "memory-use/scripts/memory.py" | |
| command = [sys.executable, str(script), "--home", str(tmp_path / "state")] | |
| result = subprocess.run(command + ["search", "query", "--json"], capture_output=True, text=True, cwd=tmp_path) | |
| assert result.returncode == 1 | |
| assert json.loads(result.stdout)["code"] == "unmapped" | |
| hook = subprocess.run(command + ["hook"], input="invalid JSON", capture_output=True, text=True) | |
| assert hook.returncode == 0 | |
| assert "systemMessage" in json.loads(hook.stdout) | |
| def test_native_network_child_has_hard_deadline(tmp_path, monkeypatch): | |
| hanging = tmp_path / "hanging.py" | |
| hanging.write_text("import time\ntime.sleep(10)\n") | |
| monkeypatch.setattr(m, "__file__", str(hanging)) | |
| with pytest.raises(m.MemoryError, match="deadline"): | |
| m.HFStore(timeout=0.05).call("whoami") | |
| def test_concurrent_saves_survive_with_distinct_objects(setup): | |
| client, store, work = setup | |
| with ThreadPoolExecutor(max_workers=4) as pool: | |
| futures = [pool.submit(client.save, f"Finding {i}", f"Concurrent {i}", cwd=work) for i in range(12)] | |
| results = [f.result() for f in futures] | |
| assert len({r["id"] for r in results}) == 12 | |
| assert all(r["status"] == "saved" for r in results) | |
| assert len(store.buckets["owner/agent-memory"]["files"]) == 12 | |
| assert not client.pending() | |
| def test_remote_traversal_skill_never_changes_installation(setup, tmp_path): | |
| client, store, work = setup | |
| setup_runtime(client, tmp_path, store, work) | |
| prefix = "skills/hostile/revision/" | |
| store.buckets["owner/agent-memory"]["files"].update({ | |
| prefix + "SKILL.md": "---\nname: hostile\ndescription: Test\n---\nInstructions", | |
| prefix + "../../../escaped": "Bad path", | |
| }) | |
| with pytest.raises(m.MemoryError, match="Unsafe"): | |
| client.install_shared("hostile", "revision", cwd=work) | |
| assert not (tmp_path / "skills/hostile").exists() | |
| assert not (tmp_path / "escaped").exists() | |
| def test_missing_authentication_has_clear_local_message(monkeypatch): | |
| import huggingface_hub | |
| monkeypatch.setattr(huggingface_hub, "get_token", lambda: None) | |
| with pytest.raises(m.MemoryError, match="No HF credentials") as result: | |
| m.hf_operation({"op": "whoami"}) | |
| assert result.value.code == "auth" | |
| def test_fresh_read_never_silently_falls_back_to_cache(setup): | |
| client, store, work = setup | |
| saved = client.save("Cached", "Fresh check", cwd=work) | |
| client.read(saved["id"], cwd=work) | |
| store.fail = "offline" | |
| assert client.read(saved["id"], cwd=work)["stale"] | |
| with pytest.raises(m.MemoryError): | |
| client.read(saved["id"], cwd=work, fresh=True) | |