Use a dataclass for pre_migration input parameters

Signed-off-by: Kartik Ganesh <[email protected]>

FetchMigration/python/pre_migration.py

@@ -7,6 +7,7 @@
 
 # Constants
 from endpoint_info import EndpointInfo
+from pre_migration_params import PreMigrationParams
 
 SUPPORTED_ENDPOINTS = ["opensearch", "elasticsearch"]
 SOURCE_KEY = "source"
@@ -157,7 +158,7 @@ def compute_endpoint_and_fetch_indices(config: dict, key: str) -> tuple[Endpoint
     return endpoint_info, indices
 
 
-def run(args: argparse.Namespace) -> None:
+def run(args: PreMigrationParams) -> None:
     # Sanity check
     if not args.report and len(args.output_file) == 0:
         raise ValueError("No output file specified")
@@ -223,4 +224,5 @@ def run(args: argparse.Namespace) -> None:
                             help="Print a report of the index differences")
     arg_parser.add_argument("--dryrun", action="store_true",
                             help="Skips the actual creation of indices on the target cluster")
-    run(arg_parser.parse_args())
+    namespace = arg_parser.parse_args()
+    run(PreMigrationParams(namespace.config_file_path, namespace.output_file, namespace.report, namespace.dryrun))

FetchMigration/python/pre_migration_params.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class PreMigrationParams:
+    config_file_path: str
+    output_file: str = ""
+    report: bool = False
+    dryrun: bool = False

FetchMigration/python/tests/test_pre_migration.py

@@ -1,4 +1,3 @@
-import argparse
 import copy
 import pickle
 import random
@@ -7,6 +6,7 @@
 from unittest.mock import patch, MagicMock, ANY
 
 import pre_migration
+from pre_migration_params import PreMigrationParams
 from tests import test_constants
 
 # Constants
@@ -251,13 +251,7 @@ def test_run_report(self, mock_fetch_indices: MagicMock, mock_create_indices: Ma
         index_settings[test_constants.INDEX_KEY][test_constants.NUM_REPLICAS_SETTING] += 1
         # Fetch indices is called first for source, then for target
         mock_fetch_indices.side_effect = [test_constants.BASE_INDICES_DATA, target_indices_data]
-        # Set up test input
-        test_input = argparse.Namespace()
-        test_input.config_file_path = test_constants.PIPELINE_CONFIG_RAW_FILE_PATH
-        # Default value for missing output file
-        test_input.output_file = ""
-        test_input.report = True
-        test_input.dryrun = False
+        test_input = PreMigrationParams(test_constants.PIPELINE_CONFIG_RAW_FILE_PATH, report=True)
         pre_migration.run(test_input)
         mock_create_indices.assert_called_once_with(expected_create_payload, ANY)
         mock_doc_count.assert_called()
@@ -280,12 +274,7 @@ def test_run_dryrun(self, mock_fetch_indices: MagicMock, mock_write_output: Magi
         del target_indices_data[index_to_create]
         # Fetch indices is called first for source, then for target
         mock_fetch_indices.side_effect = [test_constants.BASE_INDICES_DATA, target_indices_data]
-        # Set up test input
-        test_input = argparse.Namespace()
-        test_input.config_file_path = test_constants.PIPELINE_CONFIG_RAW_FILE_PATH
-        test_input.output_file = expected_output_path
-        test_input.dryrun = True
-        test_input.report = False
+        test_input = PreMigrationParams(test_constants.PIPELINE_CONFIG_RAW_FILE_PATH, expected_output_path, dryrun=True)
         pre_migration.run(test_input)
         mock_write_output.assert_called_once_with(self.loaded_pipeline_config, {index_to_create}, expected_output_path)
         mock_doc_count.assert_called()
@@ -321,12 +310,7 @@ def test_write_output(self, mock_dump: MagicMock):
             mock_dump.assert_called_once_with(expected_output_data, ANY)
 
     def test_missing_output_file_non_report(self):
-        # Set up test input
-        test_input = argparse.Namespace()
-        test_input.config_file_path = test_constants.PIPELINE_CONFIG_RAW_FILE_PATH
-        # Default value for missing output file
-        test_input.output_file = ""
-        test_input.report = False
+        test_input = PreMigrationParams(test_constants.PIPELINE_CONFIG_RAW_FILE_PATH)
         self.assertRaises(ValueError, pre_migration.run, test_input)