73 lines
2.0 KiB
Python
73 lines
2.0 KiB
Python
"""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
|