Every test module carried the same eight-line prologue redirecting the database to a temp file. Only the first one to be imported ever took effect: `app.database` reads `AIDND_DB_PATH` at import and builds `engine` from it once, so by the time the second module ran the engine already existed. The other 34 copies created a temp file that nothing opened and nothing deleted, and leaked one per module per run. `conftest.py` now does it once, which is early enough because pytest imports conftest before any test module. It also deletes the file when the run ends. The tests still share one database, exactly as they already did: each `client` fixture calls `create_all` on setup and `drop_all` on teardown, so no test sees another test's rows. `tests/fakes.py` holds the one `ScriptedProvider`. Nine modules each had a copy, and the copies had drifted into four feature sets, so a test that needed to raise a provider error had to be written in one of the files whose copy supported that. The shared one is the superset. The two `FakeProvider` copies were the same class with a fixed reply, so they use it too. `test_chat.py` keeps its own, which implements `chat` rather than `generate` and records what it was constructed with. An autouse fixture resets the fake's class state between tests, so a stale reply list can no longer reach the next test. 435 lines out of the suite. 549 tests pass. Verified live by sabotage: breaking the shared fake fails 13 tests across four modules. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014Dix4oGV3njgWRdu7P9t6r
61 lines
2.1 KiB
Python
61 lines
2.1 KiB
Python
"""Tests for the SSRF guard on the user-supplied BYOK endpoint_url.
|
|
|
|
python -m pytest tests/test_netguard.py -v
|
|
"""
|
|
import pytest
|
|
|
|
from app import auth, netguard
|
|
|
|
|
|
@pytest.fixture
|
|
def hosted(monkeypatch):
|
|
monkeypatch.setattr(auth, "MULTI_USER", True)
|
|
|
|
|
|
def _resolves_to(monkeypatch, ip: str):
|
|
"""Pin `getaddrinfo` so the test exercises the address decision, not real DNS."""
|
|
monkeypatch.setattr(
|
|
netguard.socket, "getaddrinfo",
|
|
lambda *a, **k: [(2, 1, 6, "", (ip, 443))],
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("ip", [
|
|
"127.0.0.1", # loopback
|
|
"169.254.169.254", # cloud metadata (link-local)
|
|
"10.0.0.5", # RFC1918
|
|
"192.168.1.1", # RFC1918
|
|
"172.16.0.9", # RFC1918
|
|
"0.0.0.0", # unspecified
|
|
"100.64.0.1", # carrier-grade NAT
|
|
"::1", # IPv6 loopback
|
|
"fd00::1", # IPv6 unique-local
|
|
])
|
|
def test_blocks_non_public_addresses(hosted, monkeypatch, ip):
|
|
_resolves_to(monkeypatch, ip)
|
|
assert netguard.endpoint_block_reason("https://evil.example.com/v1") is not None
|
|
|
|
|
|
def test_allows_public_address(hosted, monkeypatch):
|
|
_resolves_to(monkeypatch, "104.18.0.1") # a public IP
|
|
assert netguard.endpoint_block_reason("https://openrouter.ai/api/v1") is None
|
|
|
|
|
|
def test_rejects_non_http_scheme(hosted):
|
|
assert netguard.endpoint_block_reason("file:///etc/passwd") is not None
|
|
assert netguard.endpoint_block_reason("gopher://x/") is not None
|
|
|
|
|
|
def test_unresolvable_host_is_blocked(hosted, monkeypatch):
|
|
def boom(*a, **k):
|
|
raise netguard.socket.gaierror("no such host")
|
|
monkeypatch.setattr(netguard.socket, "getaddrinfo", boom)
|
|
assert netguard.endpoint_block_reason("https://nope.invalid/v1") is not None
|
|
|
|
|
|
def test_noop_in_local_mode(monkeypatch):
|
|
monkeypatch.setattr(auth, "MULTI_USER", False)
|
|
# Local installs must reach localhost (Ollama). The guard never blocks local mode.
|
|
assert netguard.endpoint_block_reason("http://localhost:11434/v1") is None
|
|
assert netguard.endpoint_block_reason("http://127.0.0.1:11434/v1") is None
|