fix: preserve assumptions in threat-model JSON

README.md

@@ -445,6 +445,9 @@ matching elements connecting to SQL datastores, would be `any(f.sink.oneOf(Datas
 
 With a little bit of Python code it is possible to import a threat model from JSON (notice the special format in the example found in `tests/input.json`). The following example imports the `input.json` example found in tests. Save the following code as `tm2.py`.
 
+JSON exports preserve assumptions as objects with their `name`, `description`, and
+`exclude` threat IDs, including assumptions attached to elements and excluded findings.
+
 ```python
 
 #!/usr/bin/env python3

pytm/pytm.py

@@ -532,6 +532,7 @@ def ts_tm(obj):
     return serialize(obj, nested=True)
 
 
+@to_serializable.register(Assumption)
 @to_serializable.register(Controls)
 @to_serializable.register(Data)
 @to_serializable.register(Threat)
@@ -629,6 +630,8 @@ def serialize(obj, nested=False):
 
         if isinstance(value, (Element, Data)):
             value = value.name
+        elif isinstance(value, Assumption):
+            value = to_serializable(value)
         elif hasattr(value, "model_dump") and not isinstance(value, TM):
             value = value.model_dump()
         elif isinstance(obj, Threat) and attr_name == "target":

tests/test_assumption_json.py

@@ -0,0 +1,72 @@
+"""Assumptions must survive threat-model JSON export and import."""
+
+import json
+
+import pytest
+
+from pytm import TM, Assumption, Server, loads
+from pytm.pytm import to_serializable
+
+
+@pytest.fixture(autouse=True)
+def reset_tm():
+    TM.reset()
+    yield
+    TM.reset()
+
+
+@pytest.mark.parametrize("scope", ["model", "element"])
+def test_assumption_round_trip_preserves_fields(scope):
+    assumption = Assumption(
+        "Reviewed input", exclude=["INP01", "INP02"], description="Validated upstream"
+    )
+    tm = TM("Assumptions", assumptions=[assumption] if scope == "model" else [])
+    Server("API", assumptions=[assumption] if scope == "element" else [])
+
+    output = json.dumps(tm, default=to_serializable)
+    TM.reset()
+    restored = loads(output)
+
+    assumptions = (
+        restored.assumptions if scope == "model" else restored._elements[0].assumptions
+    )
+    assert len(assumptions) == 1
+    assert assumptions[0].name == "Reviewed input"
+    assert assumptions[0].description == "Validated upstream"
+    assert assumptions[0].exclude == {"INP01", "INP02"}
+
+
+@pytest.mark.parametrize("scope", ["model", "element"])
+def test_round_trip_preserves_threat_exclusion(tmp_path, scope):
+    threats_file = tmp_path / "threats.json"
+    threats_file.write_text(
+        json.dumps(
+            [
+                {
+                    "SID": "TEST01",
+                    "target": ["Server"],
+                    "description": "Always matches the server",
+                    "condition": "True",
+                }
+            ]
+        )
+    )
+    tm = TM("Exclusions", threatsFile=str(threats_file))
+    server = Server("API")
+    tm.resolve()
+    assert [finding.threat_id for finding in tm.findings] == ["TEST01"]
+
+    owner = tm if scope == "model" else server
+    owner.assumptions = [Assumption("Reviewed", exclude=["TEST01"])]
+    tm.resolve()
+    assert tm.findings == []
+    assert [finding.threat_id for finding in tm.excluded_findings] == ["TEST01"]
+
+    output = json.dumps(tm, default=to_serializable)
+    TM.reset()
+    restored = loads(output)
+    restored.resolve()
+
+    assert restored.findings == []
+    assert [finding.threat_id for finding in restored.excluded_findings] == ["TEST01"]
+    assert restored.excluded_findings[0].assumption.exclude == {"TEST01"}