"""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