Spaces:
Running
Running
File size: 9,973 Bytes
33391b8 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 | """Unit tests for the four MCP tools, using pytest-httpx to mock the
external calls (rdap.org / IPQS / VirusTotal) since a hackathon dev loop
shouldn't burn real quota just to check that parsing logic works.
Run: pytest tests/ -v
"""
import os
from pathlib import Path
import httpx
import pytest
from mcp_server import cache, resilience
from mcp_server.blocklist import Blocklist
FIXTURES_DIR = Path(__file__).resolve().parent.parent / "data" / "blocklist_snapshot"
SAMPLE_OPENPHISH = FIXTURES_DIR / "sample_openphish.txt"
SAMPLE_URLHAUS = FIXTURES_DIR / "sample_urlhaus.csv"
@pytest.fixture(autouse=True)
def _reset_shared_state():
"""Cache and circuit-breaker state is module-level (by design — it
needs to survive across tool calls within one server process), so
each test must start from a clean slate."""
cache.cache_clear()
resilience.reset_circuit("rdap")
resilience.reset_circuit("ipqs")
resilience.reset_circuit("virustotal")
os.environ.setdefault("IPQS_API_KEY", "test-ipqs-key")
os.environ.setdefault("VIRUSTOTAL_API_KEY", "test-vt-key")
yield
# ---------------------------------------------------------------------------
# lookup_domain_age
# ---------------------------------------------------------------------------
class TestLookupDomainAge:
def test_known_good_domain_parses_correctly(self, httpx_mock):
httpx_mock.add_response(
url="https://rdap.org/domain/google.com",
json={
"events": [{"eventAction": "registration", "eventDate": "1997-09-15T04:00:00Z"}],
"entities": [{
"roles": ["registrar"],
"vcardArray": ["vcard", [["fn", {}, "text", "MarkMonitor Inc."]]],
}],
},
)
from mcp_server.server import lookup_domain_age
result = lookup_domain_age("google.com")
assert result["status"] == "ok"
assert result["registrar"] == "MarkMonitor Inc."
assert result["age_days"] > 9000 # registered 1997, comfortably old
def test_rdap_failure_returns_degraded_not_raise(self, httpx_mock):
httpx_mock.add_response(url="https://rdap.org/domain/broken.test", status_code=500)
httpx_mock.add_response(url="https://rdap.org/domain/broken.test", status_code=500)
from mcp_server.server import lookup_domain_age
result = lookup_domain_age("broken.test")
assert result["status"] == "degraded"
assert result["age_days"] is None
def test_result_is_cached_second_call_makes_no_http_request(self, httpx_mock):
httpx_mock.add_response(
url="https://rdap.org/domain/cached.test",
json={"events": [{"eventAction": "registration", "eventDate": "2020-01-01T00:00:00Z"}]},
)
from mcp_server.server import lookup_domain_age
first = lookup_domain_age("cached.test")
second = lookup_domain_age("cached.test") # no mock registered for this — would fail if it hit the network
assert first == second
assert len(httpx_mock.get_requests()) == 1
# ---------------------------------------------------------------------------
# check_ipqs
# ---------------------------------------------------------------------------
class TestCheckIpqs:
def test_known_phishing_url_returns_elevated_risk_score(self, httpx_mock):
httpx_mock.add_response(
url=httpx.URL(
"https://www.ipqualityscore.com/api/json/url/test-ipqs-key/"
"http%3A%2F%2Fpaypal-verify-account.example.net%2Fsecure"
),
json={"success": True, "risk_score": 92, "suspicious": True, "phishing": True, "malware": False},
)
from mcp_server.server import check_ipqs
result = check_ipqs("http://paypal-verify-account.example.net/secure")
assert result["status"] == "ok"
assert result["risk_score"] == 92
assert result["is_phishing"] is True
def test_known_good_url_returns_low_risk(self, httpx_mock):
httpx_mock.add_response(
url=httpx.URL("https://www.ipqualityscore.com/api/json/url/test-ipqs-key/https%3A%2F%2Fgoogle.com"),
json={"success": True, "risk_score": 1, "suspicious": False, "phishing": False, "malware": False},
)
from mcp_server.server import check_ipqs
result = check_ipqs("https://google.com")
assert result["risk_score"] == 1
assert result["is_suspicious"] is False
def test_garbage_api_key_returns_degraded(self, httpx_mock):
httpx_mock.add_response(
url=httpx.URL("https://www.ipqualityscore.com/api/json/url/test-ipqs-key/https%3A%2F%2Fexample.com"),
json={"success": False, "message": "Invalid API Key"},
)
httpx_mock.add_response(
url=httpx.URL("https://www.ipqualityscore.com/api/json/url/test-ipqs-key/https%3A%2F%2Fexample.com"),
json={"success": False, "message": "Invalid API Key"},
)
from mcp_server.server import check_ipqs
result = check_ipqs("https://example.com")
assert result["status"] == "degraded"
def test_429_first_strike_degraded_second_strike_rate_limited(self, httpx_mock):
url = httpx.URL("https://www.ipqualityscore.com/api/json/url/test-ipqs-key/https%3A%2F%2Frate.test")
from mcp_server.server import check_ipqs
httpx_mock.add_response(url=url, status_code=429)
cache.cache_clear() # ensure no cache short-circuit between calls
r1 = check_ipqs("https://rate.test")
assert r1["status"] == "degraded"
cache.cache_clear()
httpx_mock.add_response(url=url, status_code=429)
r2 = check_ipqs("https://rate.test")
assert r2["status"] == "rate_limited"
# Circuit now open: a third call shouldn't even hit the network.
cache.cache_clear()
r3 = check_ipqs("https://rate.test")
assert r3["status"] == "rate_limited"
assert len(httpx_mock.get_requests()) == 2 # not 3
# ---------------------------------------------------------------------------
# check_virustotal
# ---------------------------------------------------------------------------
class TestCheckVirustotal:
def test_known_malicious_url_reports_malicious_count(self, httpx_mock):
import base64
url = "http://malware-drop.example.org/payload.exe"
url_id = base64.urlsafe_b64encode(url.encode()).decode().strip("=")
httpx_mock.add_response(
url=f"https://www.virustotal.com/api/v3/urls/{url_id}",
json={"data": {"attributes": {"last_analysis_stats": {
"malicious": 12, "suspicious": 2, "harmless": 60, "undetected": 6,
}}}},
)
from mcp_server.server import check_virustotal
result = check_virustotal(url)
assert result["status"] == "ok"
assert result["malicious_count"] == 12
assert result["total_engines"] == 80
def test_never_before_seen_url_returns_zero_not_error(self, httpx_mock):
import base64
url = "https://brand-new-url.example.test"
url_id = base64.urlsafe_b64encode(url.encode()).decode().strip("=")
httpx_mock.add_response(url=f"https://www.virustotal.com/api/v3/urls/{url_id}", status_code=404)
from mcp_server.server import check_virustotal
result = check_virustotal(url)
assert result["status"] == "ok"
assert result["malicious_count"] == 0
def test_circuit_breaker_skips_vt_after_two_429s_in_session(self, httpx_mock):
import base64
url = "https://vt-rate-test.example.com"
url_id = base64.urlsafe_b64encode(url.encode()).decode().strip("=")
endpoint = f"https://www.virustotal.com/api/v3/urls/{url_id}"
from mcp_server.server import check_virustotal
httpx_mock.add_response(url=endpoint, status_code=429)
cache.cache_clear()
check_virustotal(url)
httpx_mock.add_response(url=endpoint, status_code=429)
cache.cache_clear()
r2 = check_virustotal(url)
assert r2["status"] == "rate_limited"
cache.cache_clear()
r3 = check_virustotal(url)
assert r3["status"] == "rate_limited"
assert len(httpx_mock.get_requests()) == 2
# ---------------------------------------------------------------------------
# check_local_blocklist (no network at all)
# ---------------------------------------------------------------------------
class TestCheckLocalBlocklist:
def test_openphish_hit(self):
bl = Blocklist(openphish_path=SAMPLE_OPENPHISH, urlhaus_path=SAMPLE_URLHAUS)
found, source = bl.contains("http://paypal-verify-account.example.net/secure")
assert found is True
assert source == "openphish"
def test_urlhaus_hit(self):
bl = Blocklist(openphish_path=SAMPLE_OPENPHISH, urlhaus_path=SAMPLE_URLHAUS)
found, source = bl.contains("http://fake-bank-login.example.io/auth")
assert found is True
assert source == "urlhaus"
def test_clean_domain_not_found(self):
bl = Blocklist(openphish_path=SAMPLE_OPENPHISH, urlhaus_path=SAMPLE_URLHAUS)
found, source = bl.contains("https://google.com")
assert found is False
assert source is None
# ---------------------------------------------------------------------------
# TTL cache unit behavior (isolated from the tools)
# ---------------------------------------------------------------------------
class TestCache:
def test_set_then_get_roundtrip(self):
cache.cache_set("k", {"a": 1})
assert cache.cache_get("k") == {"a": 1}
def test_expired_entry_returns_none(self, monkeypatch):
import time
cache.cache_set("k", {"a": 1})
future = time.time() + cache.CACHE_TTL_SECONDS + 1
monkeypatch.setattr(time, "time", lambda: future)
assert cache.cache_get("k") is None
|