Fix likelihood regression and added tests

pytm/finding.py

@@ -27,6 +27,7 @@ class Finding(BaseModel):
         assumption (Assumption): The assumption that caused this finding to be excluded
         response (str): Describes how this threat matching this particular asset or dataflow is being handled. Can be one of: mitigated, transferred, avoided, accepted
         cvss (str): The CVSS score and/or vector
+        likelihood (str): Likelihood of the threat
     """
 
     model_config = ConfigDict(
@@ -57,6 +58,7 @@ class Finding(BaseModel):
         description="Describes how this threat matching this particular asset or dataflow is being handled. Can be one of: mitigated, transferred, avoided, accepted",
     )
     cvss: str = Field(default="", description="The CVSS score and/or vector")
+    likelihood: str = Field(default="", description="Likelihood of the threat")
 
     def __init__(self, *args, **kwargs):
         """Initialize a Finding.
@@ -66,7 +68,7 @@ def __init__(self, *args, **kwargs):
             **kwargs: Finding properties:
                 - element (Element): Element this finding applies to
                 - target (str): Name of the element this finding applies to
-                - threat (Threat): Threat object to copy attributes from (description, details, severity, mitigations, example, references, condition)
+                - threat (Threat): Threat object to copy attributes from (description, details, severity, mitigations, example, references, condition, likelihood)
                 - description (str): Threat description
                 - details (str): Threat details
                 - severity (str): Threat severity
@@ -79,6 +81,7 @@ def __init__(self, *args, **kwargs):
                 - assumption (Assumption): The assumption that caused this finding to be excluded
                 - response (str): Describes how this threat matching this particular asset or dataflow is being handled. Can be one of: mitigated, transferred, avoided, accepted
                 - cvss (str): The CVSS score and/or vector
+                - likelihood (str): Likelihood of the threat
         """
         # Handle positional element argument
         if args:
@@ -105,6 +108,7 @@ def __init__(self, *args, **kwargs):
                 "example",
                 "references",
                 "condition",
+                "likelihood",
             ]
             for attr in threat_attrs:
                 if attr not in kwargs:  # Don't override explicit values

pytm/tm.py

@@ -656,6 +656,7 @@ def encode_threat_data(obj):
                 "condition",
                 "cvss",
                 "response",
+                "likelihood",
             ]
 
             items = obj if isinstance(obj, list) else [obj]

tests/test_pytmfunc.py

@@ -1728,3 +1728,44 @@ def test_override_finding_applied_through_resolve(self):
         assert len(server.findings) == 1
         assert server.findings[0].response == "accepted"
         assert server.findings[0].cvss == "5.0"
+
+    def test_likelihood_copied_from_threat_to_finding(self):
+        """likelihood is propagated from Threat to resolved Finding."""
+        TM.reset()
+        tm = TM("test tm", description="aaa")
+        Server("Web Server")
+        TM._threats = [Threat(SID="T01", target="Server", severity="High", likelihood="Medium")]
+        tm.resolve()
+
+        server = next(e for e in TM._elements if e.name == "Web Server")
+        assert len(server.findings) == 1
+        assert server.findings[0].likelihood == "Medium"
+
+    def test_override_finding_likelihood_not_overwritten(self):
+        """An explicit likelihood on a Finding override is preserved after resolve."""
+        TM.reset()
+        tm = TM("test tm", description="aaa")
+        Server(
+            "Web Server",
+            overrides=[
+                Finding(threat_id="T01", likelihood="High"),
+            ],
+        )
+        TM._threats = [Threat(SID="T01", target="Server", severity="High", likelihood="Low")]
+        tm.resolve()
+
+        server = next(e for e in TM._elements if e.name == "Web Server")
+        assert len(server.findings) == 1
+        assert server.findings[0].likelihood == "High"
+
+    def test_finding_likelihood_defaults_to_empty(self):
+        """likelihood defaults to empty string when the threat has none."""
+        TM.reset()
+        tm = TM("test tm", description="aaa")
+        Server("Web Server")
+        TM._threats = [Threat(SID="T01", target="Server", severity="High")]
+        tm.resolve()
+
+        server = next(e for e in TM._elements if e.name == "Web Server")
+        assert len(server.findings) == 1
+        assert server.findings[0].likelihood == ""