diff --git a/produce_risk_tracker.py b/produce_risk_tracker.py index 8cb74fa..e5f728f 100644 --- a/produce_risk_tracker.py +++ b/produce_risk_tracker.py @@ -17,6 +17,7 @@ def print_yaml(risk_dict: dict): def build_risk(risk_dict: dict): template_file = open("yaml-templates/risks_template.yaml") template_str = template_file.read() + template_file.close() risk_template = Template(template_str) risk_yaml = risk_template.render(risk_dict) diff --git a/tests/test_produce_risk_tracker.py b/tests/test_produce_risk_tracker.py new file mode 100644 index 0000000..8a2ed02 --- /dev/null +++ b/tests/test_produce_risk_tracker.py @@ -0,0 +1,30 @@ +import unittest + +from produce_risk_tracker import build_risk + + +class TestRiskBuilder(unittest.TestCase): + + def test_build_risk(self): + risk_dict = { + "name": "test_asset@test_test", + "status": "mitigated", + "justification": "Because security.", + "ticket": "XXX-YYY", + "date": "2024-05-16", + "checked_by": "pritchyspritch", + } + risk_yaml = build_risk(risk_dict) + + expected_yaml = """ test_asset@test_test: # wildcards "*" between the @ characters are possible + status: mitigated # values: unchecked, in-discussion, accepted, in-progress, mitigated, false-positive + justification: Because security. + ticket: XXX-YYY + date: 2024-05-16 + checked_by: pritchyspritch""" + + self.assertEqual(expected_yaml, risk_yaml) + + +if __name__ == "__main__": + unittest.main()