diff --git a/api/routes/cbom.py b/api/routes/cbom.py index 82f63e7..bb6beb7 100644 --- a/api/routes/cbom.py +++ b/api/routes/cbom.py @@ -1,5 +1,7 @@ """Cryptographic Bill of Materials routes for post-quantum findings.""" +from __future__ import annotations + import logging import os from datetime import datetime, timezone diff --git a/tests/test_api_routes.py b/tests/test_api_routes.py new file mode 100644 index 0000000..0786d3a --- /dev/null +++ b/tests/test_api_routes.py @@ -0,0 +1,324 @@ +"""Unit tests for API routes: score, compliance, drift, resources, prioritization.""" + +from unittest.mock import MagicMock, patch + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_cursor(fetchone=None, fetchall=None): + """Return a mock cursor with configurable return values.""" + cursor = MagicMock() + cursor.__enter__ = MagicMock(return_value=cursor) + cursor.__exit__ = MagicMock(return_value=False) + cursor.fetchone.return_value = fetchone + cursor.fetchall.return_value = fetchall or [] + return cursor + + +def _make_conn(fetchone=None, fetchall=None): + """Return a mock DB connection.""" + conn = MagicMock() + cursor = _make_cursor(fetchone=fetchone, fetchall=fetchall) + conn.cursor.return_value = cursor + return conn + + +def _make_db(score=85, compliance=None, fetchall=None, fetchone=None): + """Return a mock DatabaseManager.""" + db = MagicMock() + db.get_score.return_value = score + db.get_cve_summary.return_value = { + "total_findings": 0, + "exploit_count": 0, + "max_cvss_score": None, + "avg_cvss_score": None, + "cve_enrichment_status": "done", + } + db.get_compliance_score.return_value = compliance or { + "framework": "cis", + "total_controls": 10, + "passed": 8, + "failed": 2, + "score_percent": 80.0, + "controls": [], + } + # For routes that use _get_conn() directly + mock_conn = _make_conn(fetchone=fetchone, fetchall=fetchall) + db._get_conn.return_value = mock_conn + return db + + +# --------------------------------------------------------------------------- +# Score route tests +# --------------------------------------------------------------------------- + + +class TestScoreRoute: + """Tests for GET /api/score.""" + + def test_score_returns_200(self, client, auth_headers): + with patch("api.routes.score._get_db") as mock_get_db: + mock_get_db.return_value = _make_db(score=85) + resp = client.get("/api/score", headers=auth_headers) + assert resp.status_code == 200 + + def test_score_returns_score_field(self, client, auth_headers): + with patch("api.routes.score._get_db") as mock_get_db: + mock_get_db.return_value = _make_db(score=100) + resp = client.get("/api/score", headers=auth_headers) + # score route returns jsonify(int) — response is the score directly + assert resp.status_code == 200 + + def test_score_value_matches_db(self, client, auth_headers): + with patch("api.routes.score._get_db") as mock_get_db: + mock_get_db.return_value = _make_db(score=72) + resp = client.get("/api/score", headers=auth_headers) + data = resp.get_json() + assert data == 72 + + def test_score_is_100_with_perfect_score(self, client, auth_headers): + with patch("api.routes.score._get_db") as mock_get_db: + mock_get_db.return_value = _make_db(score=100) + resp = client.get("/api/score", headers=auth_headers) + data = resp.get_json() + assert data == 100 + + def test_score_returns_500_on_db_error(self, client, auth_headers): + with patch("api.routes.score._get_db", side_effect=RuntimeError("DB error")): + resp = client.get("/api/score", headers=auth_headers) + assert resp.status_code == 500 + + +# --------------------------------------------------------------------------- +# Compliance route tests +# --------------------------------------------------------------------------- + + +class TestComplianceRoute: + """Tests for GET /api/compliance/.""" + + def test_compliance_cis_returns_200(self, client, auth_headers): + with patch("api.routes.compliance._get_db") as mock_get_db: + mock_get_db.return_value = _make_db() + resp = client.get("/api/compliance/cis", headers=auth_headers) + assert resp.status_code == 200 + + def test_compliance_nist_returns_200(self, client, auth_headers): + with patch("api.routes.compliance._get_db") as mock_get_db: + mock_get_db.return_value = _make_db() + resp = client.get("/api/compliance/nist", headers=auth_headers) + assert resp.status_code == 200 + + def test_compliance_iso27001_returns_200(self, client, auth_headers): + with patch("api.routes.compliance._get_db") as mock_get_db: + mock_get_db.return_value = _make_db() + resp = client.get("/api/compliance/iso27001", headers=auth_headers) + assert resp.status_code == 200 + + def test_compliance_soc2_returns_200(self, client, auth_headers): + with patch("api.routes.compliance._get_db") as mock_get_db: + mock_get_db.return_value = _make_db() + resp = client.get("/api/compliance/soc2", headers=auth_headers) + assert resp.status_code == 200 + + def test_compliance_invalid_framework_returns_400(self, client, auth_headers): + resp = client.get("/api/compliance/unknown", headers=auth_headers) + assert resp.status_code == 400 + + def test_compliance_returns_500_on_db_error(self, client, auth_headers): + with patch("api.routes.compliance._get_db", side_effect=RuntimeError("DB error")): + resp = client.get("/api/compliance/cis", headers=auth_headers) + assert resp.status_code == 500 + + +# --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- +# Drift route tests +# --------------------------------------------------------------------------- + + +class TestDriftRoute: + """Tests for GET /api/drift.""" + + def _make_drift_db(self, scans=None): + db = MagicMock() + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_cursor.__enter__ = MagicMock(return_value=mock_cursor) + mock_cursor.__exit__ = MagicMock(return_value=False) + mock_cursor.fetchall.return_value = scans or [] + mock_conn.cursor.return_value = mock_cursor + db._get_conn.return_value = mock_conn + return db + + def test_drift_returns_200_with_no_scans(self, client, auth_headers): + with patch("api.routes.drift._get_db") as mock_get_db: + mock_get_db.return_value = self._make_drift_db(scans=[]) + resp = client.get("/api/drift", headers=auth_headers) + assert resp.status_code == 200 + + def test_drift_returns_empty_events_with_one_scan(self, client, auth_headers): + """Single scan returns empty events — no comparison possible.""" + import datetime + + scans = [{"scan_id": "scan-001", "started_at": datetime.datetime(2024, 1, 1)}] + with patch("api.routes.drift._get_db") as mock_get_db: + mock_get_db.return_value = self._make_drift_db(scans=scans) + resp = client.get("/api/drift", headers=auth_headers) + assert resp.status_code == 200 + data = resp.get_json() + assert data["events"] == [] + assert data["summary"]["total"] == 0 + + def test_drift_response_has_summary_and_events(self, client, auth_headers): + with patch("api.routes.drift._get_db") as mock_get_db: + mock_get_db.return_value = self._make_drift_db(scans=[]) + resp = client.get("/api/drift", headers=auth_headers) + data = resp.get_json() + assert "summary" in data + assert "events" in data + + def test_drift_summary_has_required_fields(self, client, auth_headers): + with patch("api.routes.drift._get_db") as mock_get_db: + mock_get_db.return_value = self._make_drift_db(scans=[]) + resp = client.get("/api/drift", headers=auth_headers) + summary = resp.get_json()["summary"] + for key in ("total", "added", "removed", "modified"): + assert key in summary + + def test_drift_returns_500_on_db_error(self, client, auth_headers): + with patch("api.routes.drift._get_db", side_effect=RuntimeError("DB error")): + resp = client.get("/api/drift", headers=auth_headers) + assert resp.status_code == 500 + + +# --------------------------------------------------------------------------- +# Resources route tests +# --------------------------------------------------------------------------- + + +class TestResourcesRoute: + """Tests for GET /api/resources.""" + + def _make_resources_db(self, fetchone=None, rows=None): + """Mock DB where fetchone() is the latest scan and fetchall() returns resources.""" + db = MagicMock() + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_cursor.__enter__ = MagicMock(return_value=mock_cursor) + mock_cursor.__exit__ = MagicMock(return_value=False) + # fetchone returns latest scan — None means no scan found + mock_cursor.fetchone.return_value = fetchone + mock_cursor.fetchall.return_value = rows or [] + # cursor_factory kwarg is ignored by MagicMock — always returns same cursor + mock_conn.cursor.return_value = mock_cursor + db._get_conn.return_value = mock_conn + return db + + def test_resources_returns_200_when_no_scan(self, client, auth_headers): + """Returns 200 with empty resources when no scan exists.""" + with patch("api.routes.resources._get_db") as mock_get_db: + mock_get_db.return_value = self._make_resources_db(fetchone=None) + resp = client.get("/api/resources", headers=auth_headers) + assert resp.status_code == 200 + + def test_resources_response_has_required_keys(self, client, auth_headers): + with patch("api.routes.resources._get_db") as mock_get_db: + mock_get_db.return_value = self._make_resources_db(fetchone=None) + resp = client.get("/api/resources", headers=auth_headers) + data = resp.get_json() + assert "resources" in data + assert "summary" in data + + def test_resources_empty_when_no_scan(self, client, auth_headers): + with patch("api.routes.resources._get_db") as mock_get_db: + mock_get_db.return_value = self._make_resources_db(fetchone=None) + resp = client.get("/api/resources", headers=auth_headers) + assert resp.get_json()["resources"] == [] + + def test_resources_returns_500_on_db_error(self, client, auth_headers): + with patch("api.routes.resources._get_db", side_effect=RuntimeError("DB error")): + resp = client.get("/api/resources", headers=auth_headers) + assert resp.status_code == 500 + + +# --------------------------------------------------------------------------- +# Prioritization route tests +# --------------------------------------------------------------------------- + + +class TestPrioritizationRoute: + """Tests for GET /api/prioritization.""" + + def _make_prioritization_db(self, fetchone=None, rules=None): + db = MagicMock() + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_cursor.__enter__ = MagicMock(return_value=mock_cursor) + mock_cursor.__exit__ = MagicMock(return_value=False) + # Single cursor is reused for all 3 execute() calls in the route: + # 1st fetchone -> scan row, fetchall -> rules, 2nd fetchone -> total count + if fetchone is None: + mock_cursor.fetchone.return_value = None + else: + mock_cursor.fetchone.side_effect = [fetchone, {"total": len(rules) if rules else 0}] + mock_cursor.fetchall.return_value = rules or [] + mock_conn.cursor.return_value = mock_cursor + db._get_conn.return_value = mock_conn + return db + + def test_prioritization_returns_200_with_no_scan(self, client, auth_headers): + with patch("api.routes.prioritization._get_db") as mock_get_db: + mock_get_db.return_value = self._make_prioritization_db(fetchone=None) + resp = client.get("/api/prioritization", headers=auth_headers) + assert resp.status_code == 200 + + def test_prioritization_response_has_required_keys(self, client, auth_headers): + with patch("api.routes.prioritization._get_db") as mock_get_db: + mock_get_db.return_value = self._make_prioritization_db(fetchone=None) + resp = client.get("/api/prioritization", headers=auth_headers) + data = resp.get_json() + assert "matrix" in data + assert "rankings" in data + assert "action_items" in data + + def test_prioritization_high_severity_ranked_first(self, client, auth_headers): + """HIGH severity rules rank before LOW severity.""" + high_rule = { + "rule_id": "AZ-NET-001", + "rule_name": "NSG broad", + "severity": "HIGH", + "category": "Network", + "affected_count": 2, + "description": "desc", + "remediation": "fix", + "resource_name": "nsg-01", + } + low_rule = { + "rule_id": "AZ-STOR-001", + "rule_name": "Tags", + "severity": "LOW", + "category": "Storage", + "affected_count": 1, + "description": "desc", + "remediation": "fix", + "resource_name": "storage-01", + } + with patch("api.routes.prioritization._get_db") as mock_get_db: + mock_get_db.return_value = self._make_prioritization_db( + fetchone={"scan_id": "s-1", "total": 3}, + rules=[high_rule, low_rule], + ) + resp = client.get("/api/prioritization", headers=auth_headers) + assert resp.status_code == 200 + rankings = resp.get_json().get("rankings", []) + if len(rankings) >= 2: + assert rankings[0]["severity"] == "HIGH" + + def test_prioritization_returns_500_on_db_error(self, client, auth_headers): + with patch("api.routes.prioritization._get_db", side_effect=RuntimeError("DB error")): + resp = client.get("/api/prioritization", headers=auth_headers) + assert resp.status_code == 500 diff --git a/tests/test_database_manager.py b/tests/test_database_manager.py new file mode 100644 index 0000000..3974ad5 --- /dev/null +++ b/tests/test_database_manager.py @@ -0,0 +1,212 @@ +"""Unit tests for DatabaseManager — connection, CRUD, and scoring logic.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from api.models.finding import DatabaseManager, Finding, SEVERITY_WEIGHTS + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_db(dsn="postgresql://test:test@localhost/test"): + """Return a DatabaseManager with a mocked psycopg2 connection.""" + with patch("api.models.finding.psycopg2.connect") as mock_connect: + db = DatabaseManager(dsn=dsn) + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_conn.cursor.return_value.__enter__ = MagicMock(return_value=mock_cursor) + mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False) + mock_connect.return_value = mock_conn + db.conn = mock_conn + return db, mock_conn, mock_cursor + + +def _make_finding(**kwargs): + defaults = { + "rule_id": "AZ-NET-001", + "rule_name": "NSG allows broad inbound", + "severity": "HIGH", + "category": "Network", + "resource_id": "/subscriptions/sub-1/rg/rg-1/nsg-1", + "resource_name": "nsg-1", + "resource_type": "Microsoft.Network/networkSecurityGroups", + "description": "NSG allows broad inbound access.", + "remediation": "Restrict inbound rules.", + "frameworks": {"CIS": "6.1"}, + "detected_at": "2024-01-01T00:00:00Z", + } + defaults.update(kwargs) + return Finding(**defaults) + + +# --------------------------------------------------------------------------- +# Initialisation tests +# --------------------------------------------------------------------------- + + +class TestDatabaseManagerInit: + """Tests for DatabaseManager initialisation.""" + + def test_init_with_explicit_dsn(self): + """DatabaseManager accepts an explicit DSN without reading environ.""" + db = DatabaseManager(dsn="postgresql://user:pass@localhost/db") + assert db.dsn == "postgresql://user:pass@localhost/db" + + def test_init_reads_database_url_from_env(self, monkeypatch): + """DatabaseManager reads DATABASE_URL when no DSN is provided.""" + monkeypatch.setenv("DATABASE_URL", "postgresql://env:env@localhost/envdb") + db = DatabaseManager() + assert db.dsn == "postgresql://env:env@localhost/envdb" + + def test_init_raises_without_dsn_or_env(self, monkeypatch): + """DatabaseManager raises KeyError when DATABASE_URL is unset.""" + monkeypatch.delenv("DATABASE_URL", raising=False) + with pytest.raises(KeyError): + DatabaseManager() + + def test_conn_is_none_before_connect(self): + """Connection is None before connect() is called.""" + db = DatabaseManager(dsn="postgresql://test:test@localhost/test") + assert db.conn is None + + +# --------------------------------------------------------------------------- +# SEVERITY_WEIGHTS tests +# --------------------------------------------------------------------------- + + +class TestSeverityWeights: + """Tests for SEVERITY_WEIGHTS constant.""" + + def test_high_outweighs_medium(self): + assert SEVERITY_WEIGHTS["HIGH"] > SEVERITY_WEIGHTS["MEDIUM"] + + def test_medium_outweighs_low(self): + assert SEVERITY_WEIGHTS["MEDIUM"] > SEVERITY_WEIGHTS["LOW"] + + def test_info_has_zero_weight(self): + assert SEVERITY_WEIGHTS.get("INFO", 0) == 0 + + def test_all_weights_non_negative(self): + for key, value in SEVERITY_WEIGHTS.items(): + assert value >= 0, f"{key} has negative weight" + + +# --------------------------------------------------------------------------- +# Finding dataclass tests +# --------------------------------------------------------------------------- + + +class TestFindingDataclass: + """Tests for the Finding dataclass.""" + + def test_finding_creation_with_required_fields(self): + """Finding can be created with required fields.""" + finding = _make_finding() + assert finding.rule_id == "AZ-NET-001" + assert finding.severity == "HIGH" + + def test_finding_to_dict_returns_dict(self): + """to_dict() returns a dictionary.""" + finding = _make_finding() + result = finding.to_dict() + assert isinstance(result, dict) + + def test_finding_to_dict_contains_rule_id(self): + """to_dict() includes rule_id.""" + finding = _make_finding() + result = finding.to_dict() + assert result["rule_id"] == "AZ-NET-001" + + def test_finding_optional_fields_default_to_none(self): + """Optional fields default to None or empty.""" + finding = _make_finding() + assert finding.scan_id is None + assert finding.playbook is None + assert finding.id is None + + def test_finding_metadata_defaults_to_empty_dict(self): + """metadata defaults to empty dict.""" + finding = _make_finding() + assert finding.metadata == {} + + def test_finding_cve_references_defaults_to_empty_list(self): + """cve_references defaults to empty list.""" + finding = _make_finding() + assert finding.cve_references == [] + + def test_finding_exploit_available_defaults_to_false(self): + """exploit_available defaults to False.""" + finding = _make_finding() + assert finding.exploit_available is False + + +# --------------------------------------------------------------------------- +# Score calculation tests +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Score calculation tests — call real get_score() via mock cursor +# --------------------------------------------------------------------------- + + +class TestScoreCalculation: + """Tests for DatabaseManager.get_score() SQL aggregation and floor logic.""" + + def _make_db_with_score(self, fetchone_value): + db = DatabaseManager(dsn="postgresql://test:test@localhost/test") + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_cursor.__enter__ = MagicMock(return_value=mock_cursor) + mock_cursor.__exit__ = MagicMock(return_value=False) + mock_cursor.fetchall.return_value = fetchone_value + mock_conn.cursor.return_value = mock_cursor + mock_conn.closed = False + db.conn = mock_conn + db._get_conn = MagicMock(return_value=mock_conn) + return db + + def test_get_score_returns_100_with_no_findings(self): + """get_score() returns 100 when there are no findings.""" + db = self._make_db_with_score([]) + assert db.get_score() == 100 + + def test_get_score_deducts_high_severity(self): + """get_score() deducts 10 per HIGH finding via SQL aggregation.""" + db = self._make_db_with_score([("HIGH", 2)]) + assert db.get_score() == 80 + + def test_get_score_deducts_multiple_severities(self): + """get_score() deducts correct amounts for mixed severities.""" + db = self._make_db_with_score([("HIGH", 1), ("MEDIUM", 2)]) + assert db.get_score() == 80 + + def test_get_score_floors_at_zero(self): + """get_score() floors at 0 — never returns negative.""" + db = self._make_db_with_score([("HIGH", 20)]) + assert db.get_score() == 0 + + def test_get_score_returns_integer(self): + """get_score() returns an integer.""" + db = self._make_db_with_score([("LOW", 1)]) + assert isinstance(db.get_score(), int) + + def test_get_score_calls_execute(self): + """get_score() calls cursor.execute — SQL actually runs.""" + db = self._make_db_with_score([]) + db.get_score() + db.conn.cursor().__enter__().execute.assert_called_once() + + def test_severity_weights_high_outweighs_medium(self): + assert SEVERITY_WEIGHTS["HIGH"] > SEVERITY_WEIGHTS["MEDIUM"] + + def test_severity_weights_medium_outweighs_low(self): + assert SEVERITY_WEIGHTS["MEDIUM"] > SEVERITY_WEIGHTS["LOW"] + + def test_severity_weights_info_is_zero(self): + assert SEVERITY_WEIGHTS.get("INFO", 0) == 0 diff --git a/tests/test_rag_pipeline.py b/tests/test_rag_pipeline.py new file mode 100644 index 0000000..846d1b9 --- /dev/null +++ b/tests/test_rag_pipeline.py @@ -0,0 +1,313 @@ +"""Unit tests for the RAG pipeline — ai/loader.py, ai/chunker.py, ai/retriever.py.""" + +from pathlib import Path +from unittest.mock import patch + +import pytest + +# --------------------------------------------------------------------------- +# ai/loader.py tests +# --------------------------------------------------------------------------- + + +class TestLoader: + """Tests for ai/loader.py document loading.""" + + def test_load_rule_documents_returns_list(self): + """load_rule_documents returns a non-empty list.""" + from ai.loader import load_rule_documents + + docs = load_rule_documents() + assert isinstance(docs, list) + assert len(docs) > 0 + + def test_load_rule_documents_have_required_keys(self): + """Each rule document has id, content and metadata keys.""" + from ai.loader import load_rule_documents + + docs = load_rule_documents() + for doc in docs: + assert "id" in doc + assert "content" in doc + assert "metadata" in doc + + def test_load_rule_documents_metadata_has_rule_id(self): + """Each rule document metadata contains rule_id.""" + from ai.loader import load_rule_documents + + docs = load_rule_documents() + for doc in docs: + assert "rule_id" in doc["metadata"] + assert doc["metadata"]["rule_id"].startswith("AZ-") + + def test_load_rule_documents_content_is_string(self): + """Document content is a non-empty string.""" + from ai.loader import load_rule_documents + + docs = load_rule_documents() + for doc in docs: + assert isinstance(doc["content"], str) + assert len(doc["content"]) > 0 + + def test_load_compliance_documents_returns_list(self): + """load_compliance_documents returns a non-empty list.""" + from ai.loader import load_compliance_documents + + docs = load_compliance_documents() + assert isinstance(docs, list) + assert len(docs) > 0 + + def test_load_compliance_documents_have_framework(self): + """Each compliance document metadata contains framework name.""" + from ai.loader import load_compliance_documents + + docs = load_compliance_documents() + known_frameworks = {"CIS Azure Benchmark", "NIST CSF", "ISO 27001", "SOC2"} + for doc in docs: + assert doc["metadata"]["framework"] in known_frameworks + + def test_load_all_documents_combines_rules_and_compliance(self): + """load_all_documents returns more docs than rules or compliance alone.""" + from ai.loader import load_all_documents, load_rule_documents, load_compliance_documents + + all_docs = load_all_documents() + rules = load_rule_documents() + compliance = load_compliance_documents() + # load_all_documents also includes skill documents in addition to rules and compliance + assert len(all_docs) >= len(rules) + len(compliance) + assert len(all_docs) > len(rules) + assert len(all_docs) > len(compliance) + + def test_load_rule_documents_ids_are_unique(self): + """Each rule document has a unique id.""" + from ai.loader import load_rule_documents + + docs = load_rule_documents() + ids = [doc["id"] for doc in docs] + assert len(ids) == len(set(ids)) + + +# --------------------------------------------------------------------------- +# ai/chunker.py tests +# --------------------------------------------------------------------------- + + +class TestChunker: + """Tests for ai/chunker.py document chunking.""" + + def _make_doc(self, content, doc_id="test-doc"): + return { + "id": doc_id, + "content": content, + "metadata": {"source": "test", "rule_id": "AZ-TEST-001"}, + } + + def test_chunk_documents_returns_list(self): + """chunk_documents returns a list.""" + from ai.chunker import chunk_documents + + docs = [self._make_doc("Short content")] + chunks = chunk_documents(docs) + assert isinstance(chunks, list) + assert len(chunks) > 0 + + def test_short_document_produces_one_chunk(self): + """A document shorter than chunk_size produces exactly one chunk.""" + from ai.chunker import chunk_documents + + docs = [self._make_doc("Short content that fits in one chunk")] + chunks = chunk_documents(docs, chunk_size=512) + assert len(chunks) == 1 + + def test_long_document_produces_multiple_chunks(self): + """A document longer than chunk_size produces multiple chunks.""" + from ai.chunker import chunk_documents + + long_content = "word " * 300 + docs = [self._make_doc(long_content)] + chunks = chunk_documents(docs, chunk_size=100, chunk_overlap=10) + assert len(chunks) > 1 + + def test_chunks_inherit_parent_metadata(self): + """Each chunk inherits metadata from its parent document.""" + from ai.chunker import chunk_documents + + docs = [self._make_doc("content", doc_id="parent-doc")] + chunks = chunk_documents(docs) + for chunk in chunks: + assert chunk["metadata"]["parent_doc_id"] == "parent-doc" + assert chunk["metadata"]["source"] == "test" + + def test_chunks_have_unique_ids(self): + """Each chunk has a unique id.""" + from ai.chunker import chunk_documents + + long_content = "word " * 300 + docs = [self._make_doc(long_content)] + chunks = chunk_documents(docs, chunk_size=100, chunk_overlap=10) + ids = [c["id"] for c in chunks] + assert len(ids) == len(set(ids)) + + def test_chunks_have_chunk_index(self): + """Each chunk metadata contains chunk_index.""" + from ai.chunker import chunk_documents + + long_content = "word " * 300 + docs = [self._make_doc(long_content)] + chunks = chunk_documents(docs, chunk_size=100, chunk_overlap=10) + for i, chunk in enumerate(chunks): + assert "chunk_index" in chunk["metadata"] + + def test_empty_documents_list(self): + """chunk_documents handles empty input gracefully.""" + from ai.chunker import chunk_documents + + chunks = chunk_documents([]) + assert chunks == [] + + def test_multiple_documents_chunked(self): + """Multiple documents are all chunked.""" + from ai.chunker import chunk_documents + + docs = [ + self._make_doc("Content A", "doc-a"), + self._make_doc("Content B", "doc-b"), + ] + chunks = chunk_documents(docs) + parent_ids = {c["metadata"]["parent_doc_id"] for c in chunks} + assert "doc-a" in parent_ids + assert "doc-b" in parent_ids + + def test_chunk_overlap_content_is_shared(self): + """Adjacent chunks share exactly the configured overlap characters.""" + from ai.chunker import chunk_documents + + alphabet = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ" + long_content = (alphabet * 4)[:200] + chunk_size = 80 + chunk_overlap = 20 + docs = [self._make_doc(long_content, "overlap-doc")] + chunks = chunk_documents(docs, chunk_size=chunk_size, chunk_overlap=chunk_overlap) + assert len(chunks) >= 2 + + tail_of_first = chunks[0]["content"][-chunk_overlap:] + head_of_second = chunks[1]["content"][:chunk_overlap] + assert tail_of_first == head_of_second, ( + f"Expected {chunk_overlap} chars of overlap.\n" + f"End of chunk 0: {repr(tail_of_first)}\n" + f"Start of chunk 1: {repr(head_of_second)}" + ) + + def test_no_overlap_when_zero(self): + """With chunk_overlap=0 the start of chunk[1] does not repeat chunk[0].""" + from ai.chunker import chunk_documents + + alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + long_content = (alphabet * 4)[:200] + chunk_size = 60 + chunk_overlap = 0 + docs = [self._make_doc(long_content, "no-overlap-doc")] + chunks = chunk_documents(docs, chunk_size=chunk_size, chunk_overlap=chunk_overlap) + assert len(chunks) >= 2 + assert chunks[0]["content"][-5:] != chunks[1]["content"][:5] + + +# --------------------------------------------------------------------------- +# ai/retriever.py tests +# --------------------------------------------------------------------------- + + +class TestRetriever: + """Tests for ai/retriever.py vector store retrieval.""" + + def test_retrieve_raises_when_chromadb_missing(self): + """retrieve raises VectorStoreNotBuilt when chromadb is unavailable.""" + from ai.retriever import retrieve, VectorStoreNotBuilt + + with patch("ai.retriever.chromadb", None): + with pytest.raises(VectorStoreNotBuilt): + retrieve("test query") + + def test_retrieve_raises_when_store_missing(self): + """retrieve raises VectorStoreNotBuilt when vectorstore dir is missing.""" + from ai.retriever import retrieve, VectorStoreNotBuilt + + with patch("ai.retriever.VECTORSTORE_DIR", Path("/nonexistent/path")): + with pytest.raises(VectorStoreNotBuilt): + retrieve("test query") + + def test_vector_store_not_built_is_runtime_error(self): + """VectorStoreNotBuilt is a subclass of RuntimeError.""" + from ai.retriever import VectorStoreNotBuilt + + assert issubclass(VectorStoreNotBuilt, RuntimeError) + + +# --------------------------------------------------------------------------- +# ai/embed.py tests +# --------------------------------------------------------------------------- + + +class TestEmbedPipeline: + """Tests for ai/embed.py build_vectorstore() pipeline.""" + + def test_build_vectorstore_raises_without_chromadb(self): + """build_vectorstore() raises RuntimeError when chromadb is unavailable.""" + import ai.embed as embed + + original = embed.chromadb + try: + embed.chromadb = None + with pytest.raises(RuntimeError, match="chromadb is not installed"): + embed.build_vectorstore() + finally: + embed.chromadb = original + + def test_build_vectorstore_raises_when_no_documents(self): + """build_vectorstore() raises RuntimeError when loader returns empty list.""" + from unittest.mock import patch + import ai.embed as embed + + if embed.chromadb is None: + pytest.skip("chromadb not installed") + with patch("ai.embed.load_all_documents", return_value=[]): + with pytest.raises(RuntimeError, match="No documents found"): + embed.build_vectorstore() + + def test_build_vectorstore_full_pipeline(self): + """build_vectorstore() creates collection, adds chunks, renames, returns count.""" + from unittest.mock import MagicMock, patch + import ai.embed as embed + + if embed.chromadb is None: + pytest.skip("chromadb not installed") + + mock_docs = [{"id": f"doc-{i}", "content": f"content {i}", "metadata": {}} for i in range(3)] + mock_chunks = [{"id": f"chunk-{i}", "content": f"chunk {i}", "metadata": {}} for i in range(5)] + + mock_collection = MagicMock() + mock_client = MagicMock() + mock_client.create_collection.return_value = mock_collection + + with patch("ai.embed.load_all_documents", return_value=mock_docs): + with patch("ai.embed.chunk_documents", return_value=mock_chunks): + with patch("ai.embed.chromadb.PersistentClient", return_value=mock_client): + with patch("ai.embed.VECTORSTORE_DIR") as mock_dir: + mock_dir.mkdir = MagicMock() + result = embed.build_vectorstore() + + # Verify collection was created with temp name + mock_client.create_collection.assert_called_once_with("openshield_temp") + + # Verify chunks were added to the collection + mock_collection.add.assert_called() + all_added_ids = [] + for c in mock_collection.add.call_args_list: + all_added_ids.extend(c.kwargs.get("ids", c.args[0] if c.args else [])) + assert len(all_added_ids) == len(mock_chunks) + + # Verify atomic rename from temp to final collection name + mock_collection.modify.assert_called_once_with(name="openshield") + + # Verify return value is chunk count + assert result == len(mock_chunks) diff --git a/tests/test_sentinel_ingest.py b/tests/test_sentinel_ingest.py new file mode 100644 index 0000000..65674b7 --- /dev/null +++ b/tests/test_sentinel_ingest.py @@ -0,0 +1,227 @@ +"""Unit tests for sentinel/ingest.py — HMAC signing and field mappings.""" + +import base64 +import hashlib +import hmac +import importlib +from unittest.mock import MagicMock, patch + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_ingest(workspace_id="test-workspace", shared_key=None): + """Import sentinel.ingest with controlled env vars.""" + import os + + if shared_key is None: + # Generate a valid base64 key + shared_key = base64.b64encode(b"test-secret-key-1234567890").decode() + with patch.dict( + os.environ, + { + "SENTINEL_WORKSPACE_ID": workspace_id, + "SENTINEL_SHARED_KEY": shared_key, + "SENTINEL_LOG_TYPE": "OpenShieldFindings", + }, + ): + import sentinel.ingest as ingest + + importlib.reload(ingest) + return ingest + + +RAW_FINDING = { + "id": "123", + "severity": "HIGH", + "detected_at": "2024-01-01T00:00:00Z", + "resource_id": "/subscriptions/sub-1/rg/rg-1/providers/Microsoft.Storage/sa", + "resource_type": "Microsoft.Storage/storageAccounts", + "resource_name": "mystorage", + "subscription_id": "sub-1", + "resource_group": "rg-1", + "region": "eastus", + "rule_id": "AZ-STOR-001", + "rule_name": "Public blob access enabled", + "description": "Storage account allows public blob access.", + "remediation": "Disable public blob access.", + "compliance": {"cis": "3.5", "nist": "PR.AC-3"}, + "tool_version": "0.2.0", +} + + +# --------------------------------------------------------------------------- +# HMAC signature tests +# --------------------------------------------------------------------------- + + +def test_build_signature_format(): + """build_signature returns a SharedKey header with correct structure.""" + ingest = _make_ingest() + date = "Mon, 01 Jan 2024 00:00:00 GMT" + sig = ingest.build_signature(date, 100) + assert sig.startswith("SharedKey test-workspace:") + assert len(sig) > 30 + + +def test_build_signature_is_deterministic(): + """Same inputs always produce the same signature.""" + ingest = _make_ingest() + date = "Mon, 01 Jan 2024 00:00:00 GMT" + sig1 = ingest.build_signature(date, 200) + sig2 = ingest.build_signature(date, 200) + assert sig1 == sig2 + + +def test_build_signature_changes_with_content_length(): + """Different content lengths produce different signatures.""" + ingest = _make_ingest() + date = "Mon, 01 Jan 2024 00:00:00 GMT" + sig1 = ingest.build_signature(date, 100) + sig2 = ingest.build_signature(date, 200) + assert sig1 != sig2 + + +def test_build_signature_changes_with_date(): + """Different dates produce different signatures.""" + ingest = _make_ingest() + sig1 = ingest.build_signature("Mon, 01 Jan 2024 00:00:00 GMT", 100) + sig2 = ingest.build_signature("Tue, 02 Jan 2024 00:00:00 GMT", 100) + assert sig1 != sig2 + + +def test_build_signature_uses_hmac_sha256(): + """Verify the HMAC-SHA256 value matches manual calculation.""" + shared_key_bytes = b"test-secret-key-1234567890" + shared_key_b64 = base64.b64encode(shared_key_bytes).decode() + ingest = _make_ingest(shared_key=shared_key_b64) + + date = "Mon, 01 Jan 2024 00:00:00 GMT" + content_length = 150 + + x_headers = f"x-ms-date:{date}" + string_to_hash = f"POST\n{content_length}\napplication/json\n{x_headers}\n/api/logs" + expected_hash = base64.b64encode( + hmac.new(shared_key_bytes, string_to_hash.encode("utf-8"), digestmod=hashlib.sha256).digest() + ).decode() + expected_sig = f"SharedKey test-workspace:{expected_hash}" + + assert ingest.build_signature(date, content_length) == expected_sig + + +# --------------------------------------------------------------------------- +# Field mapping tests +# --------------------------------------------------------------------------- + + +def test_normalise_maps_all_fields(): + """normalise maps all expected fields from a raw finding.""" + ingest = _make_ingest() + result = ingest.normalise(RAW_FINDING, "scan-001") + + assert result["ScanId"] == "scan-001" + assert result["FindingId"] == "123" + assert result["ResourceId"] == RAW_FINDING["resource_id"] + assert result["ResourceType"] == RAW_FINDING["resource_type"] + assert result["ResourceName"] == RAW_FINDING["resource_name"] + assert result["SubscriptionId"] == "sub-1" + assert result["ResourceGroup"] == "rg-1" + assert result["Region"] == "eastus" + assert result["RuleId"] == "AZ-STOR-001" + assert result["RuleName"] == "Public blob access enabled" + assert result["Severity"] == "High" + assert result["SeverityScore"] == 3 + assert result["Description"] == RAW_FINDING["description"] + assert result["Remediation"] == RAW_FINDING["remediation"] + assert result["CisControl"] == "3.5" + assert result["NistControl"] == "PR.AC-3" + assert result["Source"] == "OpenShield" + assert result["ToolVersion"] == "0.2.0" + + +def test_normalise_severity_scores(): + """SeverityScore maps correctly for all severity levels.""" + ingest = _make_ingest() + expected = {"CRITICAL": 4, "HIGH": 3, "MEDIUM": 2, "LOW": 1, "INFO": 0} + for sev, score in expected.items(): + finding = {**RAW_FINDING, "severity": sev} + result = ingest.normalise(finding, "scan-001") + assert result["SeverityScore"] == score, f"Failed for {sev}" + + +def test_normalise_unknown_severity_defaults_to_zero(): + """Unknown severity maps to SeverityScore 0.""" + ingest = _make_ingest() + finding = {**RAW_FINDING, "severity": "UNKNOWN"} + result = ingest.normalise(finding, "scan-001") + assert result["SeverityScore"] == 0 + + +def test_normalise_missing_fields_use_defaults(): + """normalise handles missing optional fields without raising.""" + ingest = _make_ingest() + result = ingest.normalise({}, "scan-002") + assert result["ScanId"] == "scan-002" + assert result["FindingId"] == "" + assert result["Source"] == "OpenShield" + assert result["Severity"] == "Medium" + + +def test_normalise_generates_timestamp_when_missing(): + """normalise generates TimeGenerated when detected_at is absent.""" + ingest = _make_ingest() + result = ingest.normalise({}, "scan-003") + assert "TimeGenerated" in result + assert result["TimeGenerated"].endswith("Z") + + +# --------------------------------------------------------------------------- +# send() tests +# --------------------------------------------------------------------------- + + +def test_send_returns_true_on_200(): + """send() returns True when the HTTP response is 200.""" + ingest = _make_ingest() + mock_response = MagicMock() + mock_response.status_code = 200 + with patch("requests.post", return_value=mock_response): + result = ingest.send([ingest.normalise(RAW_FINDING, "scan-001")]) + assert result is True + + +def test_send_returns_false_after_retries(): + """send() returns False after exactly 3 failed attempts.""" + ingest = _make_ingest() + mock_response = MagicMock() + mock_response.status_code = 500 + mock_response.text = "Internal Server Error" + with patch("requests.post", return_value=mock_response) as mock_post: + with patch("time.sleep"): + result = ingest.send([ingest.normalise(RAW_FINDING, "scan-001")]) + assert result is False + assert mock_post.call_count == 3 + + +def test_send_retries_on_exception(): + """send() makes exactly 3 attempts when requests.post raises an exception.""" + ingest = _make_ingest() + with patch("requests.post", side_effect=Exception("Connection error")) as mock_post: + with patch("time.sleep"): + result = ingest.send([ingest.normalise(RAW_FINDING, "scan-001")]) + assert result is False + assert mock_post.call_count == 3 + + +def test_send_stops_after_first_success(): + """send() returns True immediately on first success without further retries.""" + ingest = _make_ingest() + mock_response = MagicMock() + mock_response.status_code = 200 + with patch("requests.post", return_value=mock_response) as mock_post: + with patch("time.sleep"): + result = ingest.send([ingest.normalise(RAW_FINDING, "scan-001")]) + assert result is True + assert mock_post.call_count == 1