feat(compliance): data classification fields and retention policies job [FASE-3.5]

backend/tests/test_data_classification.py

@@ -0,0 +1,72 @@
+"""Tests for data classification fields and admin updates."""
+
+from app.models.enums import TestState
+from app.models.test import Test
+from app.models.technique import Technique
+
+
+def _seed_technique(db) -> Technique:
+    technique = Technique(
+        mitre_id="T9999",
+        name="Test Technique",
+        tactic="test",
+        platforms=["linux"],
+    )
+    db.add(technique)
+    db.commit()
+    db.refresh(technique)
+    return technique
+
+
+def test_new_test_defaults_to_internal(db, red_lead_user):
+    technique = _seed_technique(db)
+    test = Test(
+        technique_id=technique.id,
+        name="Classification test",
+        created_by=red_lead_user.id,
+    )
+    db.add(test)
+    db.commit()
+    db.refresh(test)
+    assert test.data_classification == "internal"
+
+
+def test_admin_can_update_classification(client, db, admin_user, admin_token, red_lead_user):
+    technique = _seed_technique(db)
+    test = Test(
+        technique_id=technique.id,
+        name="Classify me",
+        created_by=red_lead_user.id,
+        state=TestState.draft,
+    )
+    db.add(test)
+    db.commit()
+
+    response = client.patch(
+        f"/api/v1/tests/{test.id}/classification",
+        json={"data_classification": "sensitive"},
+        headers={"Authorization": f"Bearer {admin_token}"},
+    )
+    assert response.status_code == 200
+    assert response.json()["data_classification"] == "sensitive"
+
+    db.refresh(test)
+    assert test.data_classification == "sensitive"
+
+
+def test_non_admin_cannot_update_classification(client, db, admin_user, red_lead_token, red_lead_user):
+    technique = _seed_technique(db)
+    test = Test(
+        technique_id=technique.id,
+        name="Protected",
+        created_by=red_lead_user.id,
+    )
+    db.add(test)
+    db.commit()
+
+    response = client.patch(
+        f"/api/v1/tests/{test.id}/classification",
+        json={"data_classification": "restricted"},
+        headers={"Authorization": f"Bearer {red_lead_token}"},
+    )
+    assert response.status_code == 403