Spaces:
Running
Running
SooryaKiranKakkarayil
Day 2: MCP server complete - 4 tools, TTL cache, retry/circuit-breaker- lookup_domain_age, check_ipqs, check_virustotal, check_local_blocklist- Shared cache.py (1h TTL) and resilience.py (retry + circuit breaker)- Local OpenPhish/URLhaus blocklist loader with real snapshot data- Mocked pytest suite (16 tests) plus live MCP Inspector verification- Fixes: follow_redirects on RDAP bootstrap redirect; sys.path fix for direct script invocation; load_dotenv() for stdio-launched subprocess
33391b8 | """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" | |
| 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 | |