diff --git a/src/command_modules/azure-cli-batch/azure/cli/command_modules/batch/tests/test_batch_commands.py b/src/command_modules/azure-cli-batch/azure/cli/command_modules/batch/tests/test_batch_commands.py
index 4c80522a49c..76b24315b0b 100644
--- a/src/command_modules/azure-cli-batch/azure/cli/command_modules/batch/tests/test_batch_commands.py
+++ b/src/command_modules/azure-cli-batch/azure/cli/command_modules/batch/tests/test_batch_commands.py
@@ -319,7 +319,7 @@ def test_batch_format_options_name(self):
 
     def test_batch_argument_tree(self):
         # pylint: disable=too-many-statements
-        tree = _command_type.BatchArgumentTree(None)
+        tree = _command_type.BatchArgumentTree(None, None)
         self.assertEqual(list(tree), [])
 
         tree.set_request_param("pool", "azure.batch.models.PoolAddParameter")
@@ -454,38 +454,38 @@ def get_client(*args):  # pylint: disable=unused-argument
             'batch_unit_tests',
             'batch_tests_pool',
             'azure.batch.operations.pool_operations#PoolOperations.add',
-            get_client, None, None, 3, None, None)
+            get_client, None, None, 3, None, None, None)
         self.command_job = _command_type.AzureBatchDataPlaneCommand(
             'batch_unit_tests',
             'batch_tests_job',
             'azure.batch.operations.job_operations#JobOperations.add',
             get_client, None, None, 3,
-            ['job.job_manager_task', 'job.job_preparation_task', 'job.job_release_task'], None)
+            ['job.job_manager_task', 'job.job_preparation_task', 'job.job_release_task'], None, None)
         self.command_task = _command_type.AzureBatchDataPlaneCommand(
             'batch_unit_tests',
             'batch_tests_task',
             'azure.batch.operations.task_operations#TaskOperations.add',
-            get_client, None, None, 1, None, None)
+            get_client, None, None, 1, None, None, None)
         self.command_file = _command_type.AzureBatchDataPlaneCommand(
             'batch_unit_tests',
             'batch_tests_file',
             'azure.batch.operations.file_operations#FileOperations.get_from_task',
-            get_client, None, None, 3, None, None)
+            get_client, None, None, 3, None, None, None)
         self.command_list = _command_type.AzureBatchDataPlaneCommand(
             'batch_unit_tests',
             'batch_tests_list',
             'azure.batch.operations.job_operations#JobOperations.list',
-            get_client, None, None, 3, None, None)
+            get_client, None, None, 3, None, None, None)
         self.command_delete = _command_type.AzureBatchDataPlaneCommand(
             'batch_unit_tests',
             'batch_tests_delete',
             'azure.batch.operations.pool_operations#PoolOperations.delete',
-            get_client, None, None, 3, None, None)
+            get_client, None, None, 3, None, None, None)
         self.command_conflicts = _command_type.AzureBatchDataPlaneCommand(
             'batch_unit_tests',
             'batch_tests_conflicts',
             'azure.batch.operations.job_schedule_operations#JobScheduleOperations.add',
-            get_client, None, None, 4, None, None)
+            get_client, None, None, 4, None, None, None)
         return super(TestBatchLoader, self).setUp()
 
     def test_batch_build_parameters(self):