Speed up unit tests. Closes #2958 (#3289)

* Optimize tests

* Apply suggestion from @Copilot

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Apply suggestion from @Copilot

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update tests/test_crons.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Conditionally apply network related patches

---------

Co-authored-by: Matteo Lodi <30625432+mlodic@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: 3 people

tests/api_app/analyzers_manager/test_views.py

@@ -1,8 +1,11 @@
 # This file is a part of IntelOwl https://github.com/intelowlproject/IntelOwl
 # See the file 'LICENSE' for copying permission.
+from contextlib import nullcontext
 from typing import Type
 from unittest.mock import patch
 
+from django.conf import settings
+
 from api_app.analyzables_manager.models import Analyzable
 from api_app.analyzers_manager.models import AnalyzerConfig, AnalyzerReport
 from api_app.choices import Classification, PythonModuleBasePaths
@@ -29,17 +32,24 @@ def test_pull(self):
         from api_app.analyzers_manager.file_analyzers.yara_scan import YaraScan
 
         analyzer = "Yara"
-        response = self.client.post(f"{self.URL}/{analyzer}/pull")
-        self.assertEqual(response.status_code, 200)
+        ctx = (
+            patch.object(YaraScan, "update", return_value=True)
+            if settings.MOCK_CONNECTIONS
+            else nullcontext()
+        )
+        with ctx as mock_update:
+            response = self.client.post(f"{self.URL}/{analyzer}/pull")
+            self.assertEqual(response.status_code, 200)
 
-        self.client.force_authenticate(self.superuser)
+            self.client.force_authenticate(self.superuser)
 
-        with patch.object(YaraScan, "update", return_value=True):
             response = self.client.post(f"{self.URL}/{analyzer}/pull")
-        self.assertEqual(response.status_code, 200, response.json())
-        result = response.json()
-        self.assertIn("status", result)
-        self.assertTrue(result["status"])
+            self.assertEqual(response.status_code, 200, response.json())
+            result = response.json()
+            self.assertIn("status", result)
+            self.assertTrue(result["status"])
+            if mock_update is not None:
+                self.assertEqual(mock_update.call_count, 2)
 
         analyzer = "Doc_Info"
         response = self.client.post(f"{self.URL}/{analyzer}/pull")

tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_capa_info.py

@@ -1,6 +1,8 @@
 import subprocess
 from unittest.mock import MagicMock, patch
 
+from django.conf import settings
+
 from api_app.analyzers_manager.file_analyzers.capa_info import CapaInfo
 
 from .base_test_class import BaseFileAnalyzerTest
@@ -29,7 +31,7 @@ def get_mocked_response(self):
         mock_requests_get = MagicMock()
         mock_requests_get.json.return_value = {"tag_name": "v1.0.0"}
 
-        return [
+        patches = [
             patch.object(CapaInfo, "update", return_value=True),
             patch("subprocess.run", return_value=response_from_command),
             patch(
@@ -39,6 +41,10 @@ def get_mocked_response(self):
             patch.object(CapaInfo, "_check_if_latest_version", return_value=True),
         ]
 
+        if settings.MOCK_CONNECTIONS:
+            patches.insert(1, patch.object(CapaInfo, "_download_signatures", return_value=None))
+        return patches
+
     def get_extra_config(self):
         return {
             "shellcode": False,

tests/api_app/analyzers_manager/unit_tests/file_analyzers/test_virushee.py

@@ -1,6 +1,7 @@
 from unittest.mock import patch
 
 import requests
+from django.conf import settings
 
 from api_app.analyzers_manager.file_analyzers.virushee import VirusheeFileUpload
 
@@ -18,7 +19,7 @@ def get_extra_config(self):
         }
 
     def get_mocked_response(self):
-        return [
+        patches = [
             patch(
                 "requests.Session.get",
                 side_effect=[
@@ -36,6 +37,10 @@ def get_mocked_response(self):
             ),
         ]
 
+        if settings.MOCK_CONNECTIONS:
+            patches.append(patch("time.sleep", return_value=None))
+        return patches
+
     class MockUpResponse:
         """Simple mock response class to simulate requests.Response"""
 

tests/api_app/analyzers_manager/unit_tests/observable_analyzers/test_pulsedive.py

@@ -1,5 +1,7 @@
 from unittest.mock import patch
 
+from django.conf import settings
+
 from api_app.analyzers_manager.observable_analyzers.pulsedive import Pulsedive
 from tests.api_app.analyzers_manager.unit_tests.observable_analyzers.base_test_class import (
     BaseAnalyzerTest,
@@ -12,7 +14,7 @@ class PulsediveTestCase(BaseAnalyzerTest):
 
     @staticmethod
     def get_mocked_response():
-        return [
+        patches = [
             patch(
                 "requests.get",
                 side_effect=[
@@ -25,6 +27,10 @@ def get_mocked_response():
             patch("requests.post", return_value=MockUpResponse({"qid": 1}, 200)),
         ]
 
+        if settings.MOCK_CONNECTIONS:
+            patches.append(patch("time.sleep", return_value=None))
+        return patches
+
     @classmethod
     def get_extra_config(cls) -> dict:
         return {"scan_mode": "active", "_api_key_name": "test_api_key", "probe": 1}

tests/api_app/test_api.py

@@ -4,7 +4,9 @@
 import datetime
 import hashlib
 import os
+from contextlib import nullcontext
 from typing import Tuple
+from unittest.mock import patch
 
 from django.conf import settings
 from django.contrib.auth import get_user_model
@@ -16,6 +18,7 @@
 from api_app.choices import Classification
 from api_app.connectors_manager.models import ConnectorConfig
 from api_app.playbooks_manager.models import PlaybookConfig
+from tests.mock_utils import if_mock_connections
 
 from .. import CustomViewSetTestCase
 
@@ -101,7 +104,8 @@ def test_ask_analysis_availability__run_all_analyzers(self):
         response = self.client.post("/api/ask_analysis_availability", data, format="json")
         self.assertEqual(response.status_code, 200)
 
-    def test_analyze_file__pcap(self):
+    @if_mock_connections(patch("intel_owl.tasks.job_pipeline.apply_async"))
+    def test_analyze_file__pcap(self, mock_apply_async=None):
         # set a fake API key or YARAify_File_Scan will be skipped as not configured
         models.PluginConfig.objects.create(
             owner=self.user,
@@ -144,6 +148,10 @@ def test_analyze_file__pcap(self):
             list(job.analyzers_to_execute.all().values_list("name", flat=True)),
         )
 
+        if mock_apply_async is not None:
+            mock_apply_async.assert_called_once()
+            self.assertEqual(mock_apply_async.call_args.kwargs["args"], [job_id])
+
     def test_analyze_file__exe(self):
         data = self.analyze_file_data.copy()
         response = self.client.post("/api/analyze_file", data, format="multipart")
@@ -232,7 +240,8 @@ def test_analyze_observable__ip(self):
         self.assertEqual(data["observable_classification"], job.analyzable.classification, msg=msg)
         self.assertEqual(self.observable_md5, job.analyzable.md5, msg=msg)
 
-    def test_analyze_observable__guess_optional(self):
+    @if_mock_connections(patch("intel_owl.tasks.job_pipeline.apply_async"))
+    def test_analyze_observable__guess_optional(self, mock_apply_async=None):
         data = self.analyze_observable_ip_data.copy()
         observable_classification = data.pop("observable_classification")  # let the server calc it
 
@@ -252,16 +261,22 @@ def test_analyze_observable__guess_optional(self):
         self.assertEqual(observable_classification, job.analyzable.classification, msg=msg)
         self.assertEqual(self.observable_md5, job.analyzable.md5, msg=msg)
 
-    def test_analyze_multiple_observables(self):
+        if mock_apply_async is not None:
+            mock_apply_async.assert_called_once()
+            self.assertEqual(mock_apply_async.call_args.kwargs["args"], [job_id])
+
+    @if_mock_connections(patch("intel_owl.tasks.job_pipeline.apply_async"))
+    def test_analyze_multiple_observables(self, mock_apply_async=None):
         data = self.mixed_observable_data.copy()
 
         response = self.client.post("/api/analyze_multiple_observables", data, format="json")
         contents = response.json()
         msg = (response.status_code, contents)
         self.assertEqual(response.status_code, 200, msg=msg)
+        if mock_apply_async is not None:
+            self.assertEqual(mock_apply_async.call_count, len(data["observables"]))
 
         content = contents["results"][0]
-
         job_id = int(content["job_id"])
         job = models.Job.objects.get(pk=job_id)
         self.assertEqual(data["observables"][0][1], job.analyzable.name, msg=msg)
@@ -275,9 +290,10 @@ def test_analyze_multiple_observables(self):
             list(job.analyzers_to_execute.all().values_list("name", flat=True)),
             msg=msg,
         )
+        if mock_apply_async is not None:
+            self.assertEqual(mock_apply_async.call_args_list[0].kwargs["args"], [job_id])
 
         content = contents["results"][1]
-
         job_id = int(content["job_id"])
         job = models.Job.objects.get(pk=job_id)
         self.assertEqual(data["observables"][1][1], job.analyzable.name, msg=msg)
@@ -286,6 +302,8 @@ def test_analyze_multiple_observables(self):
             list(job.analyzers_to_execute.all().values_list("name", flat=True)),
             msg=msg,
         )
+        if mock_apply_async is not None:
+            self.assertEqual(mock_apply_async.call_args_list[1].kwargs["args"], [job_id])
         job.delete()
 
     def test_observable_no_analyzers_only_connector(self):
@@ -527,7 +545,12 @@ def test_job_rescan__observable_playbook(self):
                 "visualizers": {},
             },
         )
-        response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
+        # dont actually run the analyzers when mocking connections, they are tested in unit tests
+        ctx = (
+            patch("intel_owl.tasks.job_pipeline.apply_async") if settings.MOCK_CONNECTIONS else nullcontext()
+        )
+        with ctx as mock_apply_async:
+            response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
         contents = response.json()
         self.assertEqual(response.status_code, 202, contents)
         new_job_id = int(contents["id"])
@@ -543,6 +566,9 @@ def test_job_rescan__observable_playbook(self):
                 "visualizers": {},
             },
         )
+        if mock_apply_async is not None:
+            mock_apply_async.assert_called_once()
+            self.assertEqual(mock_apply_async.call_args.kwargs["args"], [new_job_id])
         an.delete()
 
     def test_job_rescan__sample_analyzers(self):
@@ -567,7 +593,11 @@ def test_job_rescan__sample_analyzers(self):
         )
         job.analyzers_requested.set([AnalyzerConfig.objects.get(name="Strings_Info")])
 
-        response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
+        ctx = (
+            patch("intel_owl.tasks.job_pipeline.apply_async") if settings.MOCK_CONNECTIONS else nullcontext()
+        )
+        with ctx as mock_apply_async:
+            response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
         contents = response.json()
         self.assertEqual(response.status_code, 202, contents)
         new_job_id = int(contents["id"])
@@ -579,6 +609,9 @@ def test_job_rescan__sample_analyzers(self):
             list(new_job.analyzers_requested.all()),
             [AnalyzerConfig.objects.get(name="Strings_Info")],
         )
+        if mock_apply_async is not None:
+            mock_apply_async.assert_called_once()
+            self.assertEqual(mock_apply_async.call_args.kwargs["args"], [new_job_id])
         self.assertEqual(
             new_job.runtime_configuration,
             {
@@ -617,7 +650,11 @@ def test_job_rescan__sample_playbook(self):
             },
         )
 
-        response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
+        ctx = (
+            patch("intel_owl.tasks.job_pipeline.apply_async") if settings.MOCK_CONNECTIONS else nullcontext()
+        )
+        with ctx as mock_apply_async:
+            response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
         contents = response.json()
         self.assertEqual(response.status_code, 202, contents)
         new_job_id = int(contents["id"])
@@ -642,6 +679,9 @@ def test_job_rescan__sample_playbook(self):
                 "visualizers": {},
             },
         )
+        if mock_apply_async is not None:
+            mock_apply_async.assert_called_once()
+            self.assertEqual(mock_apply_async.call_args.kwargs["args"], [new_job_id])
         job.delete()
         an.delete()
 
@@ -665,14 +705,24 @@ def test_job_rescan__permission(self):
                 "visualizers": {},
             },
         )
-        # same user
-        response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
-        contents = response.json()
-        self.assertEqual(response.status_code, 202, contents)
-        # another user
-        self.client.logout()
-        self.client.force_login(self.guest)
-        response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
-        contents = response.json()
-        self.assertEqual(response.status_code, 403, contents)
+        ctx = (
+            patch("intel_owl.tasks.job_pipeline.apply_async") if settings.MOCK_CONNECTIONS else nullcontext()
+        )
+        with ctx as mock_apply_async:
+            response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
+            contents = response.json()
+            self.assertEqual(response.status_code, 202, contents)
+            new_job_id = int(contents["id"])
+            if mock_apply_async is not None:
+                mock_apply_async.assert_called_once()
+                self.assertEqual(mock_apply_async.call_args.kwargs["args"], [new_job_id])
+                mock_apply_async.reset_mock()
+
+            self.client.logout()
+            self.client.force_login(self.guest)
+            response = self.client.post(f"/api/jobs/{job.pk}/rescan", format="json")
+            contents = response.json()
+            self.assertEqual(response.status_code, 403, contents)
+            if mock_apply_async is not None:
+                mock_apply_async.assert_not_called()
         an.delete()

tests/test_crons.py

@@ -242,9 +242,34 @@ def test_quark_updater(self):
         quark_engine.QuarkEngine.update()
         self.assertTrue(os.path.exists(DIR_PATH))
 
-    def test_yara_updater(self):
-        yara_scan.YaraScan.update()
-        self.assertTrue(len(os.listdir(settings.YARA_RULES_PATH)))
+    @if_mock_connections(
+        patch("git.Repo"),
+        patch("requests.get", return_value=MockUpResponse({}, 200)),
+        patch("zipfile.ZipFile"),
+    )
+    def test_yara_updater(self, mock_zipfile=None, mock_get=None, mock_repo=None):
+        if mock_zipfile is None or mock_get is None or mock_repo is None:
+            yara_scan.YaraScan.update()
+            self.assertTrue(os.path.isdir(settings.YARA_RULES_PATH))
+        else:
+
+            def create_yara_file(path):
+                os.makedirs(path, exist_ok=True)
+                yara_file = os.path.join(path, "test_rule.yar")
+                with open(yara_file, "w") as f:
+                    f.write(
+                        "rule TestRule {\n"
+                        "    strings:\n"
+                        '        $test = "test"\n'
+                        "    condition:\n"
+                        "        $test\n"
+                        "}\n"
+                    )
+
+            mock_repo.clone_from.side_effect = lambda url, path, **kwargs: create_yara_file(path)
+            mock_zipfile.return_value.extractall.side_effect = create_yara_file
+            result = yara_scan.YaraScan.update()
+            self.assertTrue(result)
 
     @if_mock_connections(
         patch(