Refactor code for future numpy removal (#946)

nettacker/core/app.py

@@ -6,7 +6,6 @@
 from threading import Thread
 
 import multiprocess
-import numpy
 
 from nettacker import logger
 from nettacker.config import Config, version_info
@@ -26,7 +25,7 @@
 from nettacker.core.messages import messages as _
 from nettacker.core.module import Module
 from nettacker.core.socks_proxy import set_socks_proxy
-from nettacker.core.utils import common as utils
+from nettacker.core.utils import common as common_utils
 from nettacker.core.utils.common import wait_for_threads_to_finish
 from nettacker.database.db import find_events, remove_old_logs
 from nettacker.database.mysql import mysql_create_database, mysql_create_tables
@@ -192,7 +191,7 @@ def run(self):
         Returns:
             True when it ends
         """
-        scan_id = utils.generate_random_token(32)
+        scan_id = common_utils.generate_random_token(32)
         log.info("ScanID: {0}".format(scan_id))
         log.info(_("regrouping_targets"))
         # find total number of targets + types + expand (subdomain, IPRanges, etc)
@@ -210,18 +209,9 @@ def run(self):
         return exit_code
 
     def start_scan(self, scan_id):
-        number_of_total_targets = len(self.arguments.targets)
-        target_groups = [
-            targets.tolist()
-            for targets in numpy.array_split(
-                self.arguments.targets,
-                (
-                    self.arguments.set_hardware_usage
-                    if self.arguments.set_hardware_usage <= len(self.arguments.targets)
-                    else number_of_total_targets
-                ),
-            )
-        ]
+        target_groups = common_utils.generate_target_groups(
+            self.arguments.targets, self.arguments.set_hardware_usage
+        )
         log.info(_("removing_old_db_records"))
 
         for target_group in target_groups:
@@ -239,7 +229,7 @@ def start_scan(self, scan_id):
         for _i in range(target_groups.count([])):
             target_groups.remove([])
 
-        log.info(_("start_multi_process").format(number_of_total_targets, len(target_groups)))
+        log.info(_("start_multi_process").format(len(self.arguments.targets), len(target_groups)))
         active_processes = []
         for t_id, target_groups in enumerate(target_groups):
             process = multiprocess.Process(

nettacker/core/arg_parser.py

@@ -17,7 +17,7 @@
 )
 from nettacker.core.messages import messages as _
 from nettacker.core.template import TemplateLoader
-from nettacker.core.utils import common as utils
+from nettacker.core.utils import common as common_utils
 from nettacker.logger import TerminalCodes, get_logger
 
 log = get_logger()
@@ -92,7 +92,7 @@ def load_modules(limit=-1, full_details=False):
             if len(module_names) == limit:
                 module_names["..."] = {}
                 break
-        module_names = utils.sort_dictionary(module_names)
+        module_names = common_utils.sort_dictionary(module_names)
         module_names["all"] = {}
 
         return module_names
@@ -118,11 +118,11 @@ def load_profiles(limit=-1):
                 else:
                     profiles[tag].append(key)
                 if len(profiles) == limit:
-                    profiles = utils.sort_dictionary(profiles)
+                    profiles = common_utils.sort_dictionary(profiles)
                     profiles["..."] = []
                     profiles["all"] = []
                     return profiles
-        profiles = utils.sort_dictionary(profiles)
+        profiles = common_utils.sort_dictionary(profiles)
         profiles["all"] = []
 
         return profiles
@@ -612,7 +612,9 @@ def parse_arguments(self):
         # threading & processing
         if options.set_hardware_usage not in {"low", "normal", "high", "maximum"}:
             die_failure(_("wrong_hardware_usage"))
-        options.set_hardware_usage = utils.select_maximum_cpu_core(options.set_hardware_usage)
+        options.set_hardware_usage = common_utils.select_maximum_cpu_core(
+            options.set_hardware_usage
+        )
 
         options.thread_per_host = int(options.thread_per_host)
         if options.thread_per_host < 1:

nettacker/core/utils/common.py

@@ -11,6 +11,8 @@
 import sys
 import time
 
+import numpy
+
 from nettacker import logger
 
 log = logger.get_logger()
@@ -181,9 +183,21 @@ def generate_and_replace_md5(content):
     )
 
 
-def arrays_to_matrix(arrays):
-    import numpy
+def generate_target_groups(targets, set_hardware_usage):
+    if not targets:
+        return targets
+
+    targets_total = len(targets)
+    return [
+        targets.tolist()
+        for targets in numpy.array_split(
+            targets,
+            (set_hardware_usage if set_hardware_usage <= targets_total else targets_total),
+        )
+    ]
+
 
+def arrays_to_matrix(arrays):
     return (
         numpy.array(numpy.meshgrid(*[arrays[array_name] for array_name in arrays]))
         .T.reshape(-1, len(arrays.keys()))

tests/core/utils/test_common.py

@@ -0,0 +1,107 @@
+from unittest.mock import patch
+
+from nettacker.core.utils import common as common_utils
+from tests.common import TestCase
+
+
+class TestCommon(TestCase):
+    def test_arrays_to_matrix(self):
+        (
+            self.assertEqual(
+                sorted(
+                    common_utils.arrays_to_matrix(
+                        {"ports": [1, 2, 3, 4, 5]},
+                    )
+                ),
+                [[1], [2], [3], [4], [5]],
+            ),
+        )
+
+        self.assertEqual(
+            sorted(
+                common_utils.arrays_to_matrix(
+                    {"x": [1, 2], "y": [3, 4], "z": [5, 6]},
+                )
+            ),
+            [
+                [1, 3, 5],
+                [1, 3, 6],
+                [1, 4, 5],
+                [1, 4, 6],
+                [2, 3, 5],
+                [2, 3, 6],
+                [2, 4, 5],
+                [2, 4, 6],
+            ],
+        )
+
+    def test_generate_target_groups_empty_list(self):
+        targets = []
+        set_hardware_usage = 3
+        result = common_utils.generate_target_groups(targets, set_hardware_usage)
+        assert result == []
+
+    def test_generate_target_groups_set_hardware_less_than_targets_total(self):
+        targets = [1, 2, 3, 4, 5]
+        set_hardware_usage = 2
+        result = common_utils.generate_target_groups(targets, set_hardware_usage)
+        assert result == [[1, 2, 3], [4, 5]]
+
+    def test_generate_target_groups_set_hardware_equal_to_targets_total(self):
+        targets = [1, 2, 3, 4, 5]
+        set_hardware_usage = 5
+        result = common_utils.generate_target_groups(targets, set_hardware_usage)
+        assert result == [[1], [2], [3], [4], [5]]
+
+    def test_generate_target_groups_set_hardware_greater_than_targets_total(self):
+        targets = [1, 2, 3]
+        set_hardware_usage = 5
+        result = common_utils.generate_target_groups(targets, set_hardware_usage)
+        assert result == [[1], [2], [3]]
+
+    def test_sort_dictionary(self):
+        input_dict = {
+            "a": 1,
+            "c": 3,
+            "d": 23,
+            "b": 2,
+        }
+        expected_dict = {
+            "a": 1,
+            "b": 2,
+            "c": 3,
+            "d": 23,
+        }
+        input_dict_keys = tuple(input_dict.keys())
+        expected_dict_keys = tuple(expected_dict.keys())
+        self.assertNotEqual(input_dict_keys, expected_dict_keys)
+
+        sorted_dict_keys = tuple(common_utils.sort_dictionary(input_dict).keys())
+        self.assertEqual(sorted_dict_keys, expected_dict_keys)
+
+    @patch("multiprocessing.cpu_count")
+    def test_select_maximum_cpu_core(self, cpu_count_mock):
+        cores_mapping = {
+            1: {"low": 1, "normal": 1, "high": 1, "maximum": 1},
+            2: {"low": 1, "normal": 1, "high": 1, "maximum": 1},
+            4: {"low": 1, "normal": 1, "high": 2, "maximum": 3},
+            6: {"low": 1, "normal": 1, "high": 3, "maximum": 5},
+            8: {"low": 1, "normal": 2, "high": 4, "maximum": 7},
+            10: {"low": 1, "normal": 2, "high": 5, "maximum": 9},
+            12: {"low": 1, "normal": 3, "high": 6, "maximum": 11},
+            16: {"low": 2, "normal": 4, "high": 8, "maximum": 15},
+            32: {"low": 4, "normal": 8, "high": 16, "maximum": 31},
+            48: {"low": 6, "normal": 12, "high": 24, "maximum": 47},
+            64: {"low": 8, "normal": 16, "high": 32, "maximum": 63},
+        }
+        for num_cores, levels in cores_mapping.items():
+            cpu_count_mock.return_value = num_cores
+            for level in ("low", "normal", "high", "maximum"):
+                self.assertEqual(
+                    common_utils.select_maximum_cpu_core(level),
+                    levels[level],
+                    f"It should be {common_utils.select_maximum_cpu_core(level)} "
+                    "of {num_cores} cores for '{level}' mode",
+                )
+
+            self.assertEqual(common_utils.select_maximum_cpu_core("invalid"), 1)