Add test for on_failure mechanism

Signed-off-by: Keshav Priyadarshi <[email protected]>

Author: keshav-space

vulnerabilities/tests/pipelines/test_base_pipeline.py

@@ -19,6 +19,8 @@
 from vulnerabilities.importer import AffectedPackage
 from vulnerabilities.importer import Reference
 from vulnerabilities.pipelines import VulnerableCodeBaseImporterPipeline
+from vulnerabilities.pipelines import VulnerableCodePipeline
+from vulnerabilities.tests.pipelines import TestLogger
 
 advisory_data1 = AdvisoryData(
     aliases=["CVE-2020-13371337"],
@@ -47,6 +49,33 @@ def get_advisory1(created_by="test_pipeline"):
     )
 
 
+class TestVulnerableCodePipeline(TestCase):
+    def test_on_failure(self):
+        class TestPipeline(VulnerableCodePipeline):
+            def __init__(self, test_logger):
+                super().__init__()
+                self.log = test_logger.write
+
+            @classmethod
+            def steps(cls):
+                return (cls.step1,)
+
+            def step1(self):
+                raise Exception("Something went wrong!")
+
+            def on_failure(self):
+                self.log("Doing cleanup.")
+
+        logger = TestLogger()
+        pipeline = TestPipeline(test_logger=logger)
+
+        pipeline.execute()
+        log_result = logger.getvalue()
+
+        self.assertIn("Pipeline failed", log_result)
+        self.assertIn("Running [on_failure] tasks", log_result)
+
+
 class TestVulnerableCodeBaseImporterPipeline(TestCase):
     @patch.object(
         VulnerableCodeBaseImporterPipeline,