diff --git a/src/fleet/HISTORY.rst b/src/fleet/HISTORY.rst index 5b9d2f46e0a..27e8e8e2223 100644 --- a/src/fleet/HISTORY.rst +++ b/src/fleet/HISTORY.rst @@ -133,4 +133,8 @@ Release History * Upgrade SDK version to 2025-04-01-preview * Add Fleet Gates support * Add TargetKubernetesVersion channel support -* Add Fleet Member labels support \ No newline at end of file +* Add Fleet Member labels support + +1.6.1 +++++++ +* Modified parameter handling to accept both file paths and inline JSON strings for the --stages argument diff --git a/src/fleet/azext_fleet/_params.py b/src/fleet/azext_fleet/_params.py index fde6ffa6c2a..1e2eab45846 100644 --- a/src/fleet/azext_fleet/_params.py +++ b/src/fleet/azext_fleet/_params.py @@ -95,8 +95,8 @@ def load_arguments(self, _): c.argument('kubernetes_version', validator=validate_kubernetes_version) c.argument('node_image_selection', arg_type=get_enum_type(['Latest', 'Consistent']), help='Node Image Selection is an option that lets you choose how your clusters\' nodes are upgraded') - c.argument('stages', type=file_type, completer=FilesCompleter(), - help='Path to a json file that defines stages to upgrade a fleet. See examples for further reference.') + c.argument('stages', + help='Path to a JSON file that defines stages to upgrade a fleet, or a JSON string. See examples for further reference.') c.argument('update_strategy_name', validator=validate_update_strategy_name, help='The name of the update strategy to use for this update run. If not specified, the default update strategy will be used.') @@ -110,8 +110,8 @@ def load_arguments(self, _): c.argument('fleet_name', options_list=['--fleet-name', '-f'], help='Specify the fleet name.') with self.argument_context('fleet updatestrategy create') as c: - c.argument('stages', type=file_type, completer=FilesCompleter(), - help='Path to a json file that defines an update strategy.') + c.argument('stages', + help='Path to a JSON file that defines an update strategy, or a JSON string.') with self.argument_context('fleet autoupgradeprofile') as c: c.argument('name', options_list=['--name', '-n'], help='Specify name for the auto upgrade profile.') diff --git a/src/fleet/azext_fleet/custom.py b/src/fleet/azext_fleet/custom.py index 96207b66203..1a9df1eb632 100644 --- a/src/fleet/azext_fleet/custom.py +++ b/src/fleet/azext_fleet/custom.py @@ -4,12 +4,11 @@ # -------------------------------------------------------------------------------------------- import os -import json from knack.util import CLIError from azure.cli.core.commands.client_factory import get_subscription_id -from azure.cli.core.util import sdk_no_wait +from azure.cli.core.util import sdk_no_wait, get_file_json, shell_safe_json_parse from azext_fleet._client_factory import CUSTOM_MGMT_FLEET from azext_fleet._helpers import print_or_merge_credentials @@ -462,9 +461,11 @@ def get_update_run_strategy(cmd, operation_group, stages): if stages is None: return None - with open(stages, 'r', encoding='utf-8') as fp: - data = json.load(fp) - fp.close() + # Check if the input is a file path or inline JSON + if os.path.exists(stages): + data = get_file_json(stages) + else: + data = shell_safe_json_parse(stages) update_group_model = cmd.get_models( "UpdateGroup", diff --git a/src/fleet/azext_fleet/tests/latest/test_stages_json.py b/src/fleet/azext_fleet/tests/latest/test_stages_json.py new file mode 100644 index 00000000000..bbbe185035c --- /dev/null +++ b/src/fleet/azext_fleet/tests/latest/test_stages_json.py @@ -0,0 +1,192 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import unittest +import json +import tempfile +import os +from unittest.mock import MagicMock, patch +from azext_fleet.custom import get_update_run_strategy +from azext_fleet.vendored_sdks.v2025_04_01_preview.models import UpdateRunStrategy, UpdateStage, UpdateGroup +from azure.cli.core.azclierror import ( + InvalidArgumentValueError, +) + +class TestStagesJsonHandling(unittest.TestCase): + """Test inline JSON support for --stages argument in fleet commands.""" + + def setUp(self): + """Set up test data and mock objects.""" + self.test_data = { + "stages": [ + { + "name": "stage1", + "groups": [ + {"name": "group1"}, + {"name": "group2"} + ], + "afterStageWaitInSeconds": 3600 + } + ] + } + + # Mock cmd object that provides get_models method + self.mock_cmd = MagicMock() + + # Set up get_models to return our mock classes + def mock_get_models(model_name, **kwargs): + if model_name == "UpdateGroup": + return UpdateGroup + elif model_name == "UpdateStage": + return UpdateStage + elif model_name == "UpdateRunStrategy": + return UpdateRunStrategy + else: + return MagicMock() + + self.mock_cmd.get_models = mock_get_models + + def test_file_path_stages(self): + """Test that file paths for stages work correctly with get_update_run_strategy.""" + # Create a temporary file with test data + with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f: + json.dump(self.test_data, f) + temp_file_path = f.name + + try: + # Test the actual function + result = get_update_run_strategy(self.mock_cmd, "fleet_update_runs", temp_file_path) + + # Verify the returned strategy + self.assertIsNotNone(result) + self.assertIsInstance(result, UpdateRunStrategy) + self.assertEqual(len(result.stages), 1) + + # Verify first stage + stage = result.stages[0] + self.assertIsInstance(stage, UpdateStage) + self.assertEqual(stage.name, "stage1") + self.assertEqual(stage.after_stage_wait_in_seconds, 3600) + self.assertEqual(len(stage.groups), 2) + + # Verify groups + self.assertIsInstance(stage.groups[0], UpdateGroup) + self.assertEqual(stage.groups[0].name, "group1") + self.assertEqual(stage.groups[1].name, "group2") + + finally: + # Clean up + os.unlink(temp_file_path) + + def test_inline_json_stages(self): + """Test that inline JSON strings work correctly with get_update_run_strategy.""" + inline_json = json.dumps(self.test_data) + + # Test the actual function + result = get_update_run_strategy(self.mock_cmd, "fleet_update_runs", inline_json) + + # Verify the returned strategy + self.assertIsNotNone(result) + self.assertIsInstance(result, UpdateRunStrategy) + self.assertEqual(len(result.stages), 1) + + # Verify first stage + stage = result.stages[0] + self.assertIsInstance(stage, UpdateStage) + self.assertEqual(stage.name, "stage1") + self.assertEqual(stage.after_stage_wait_in_seconds, 3600) + self.assertEqual(len(stage.groups), 2) + + # Verify groups + self.assertIsInstance(stage.groups[0], UpdateGroup) + self.assertEqual(stage.groups[0].name, "group1") + self.assertEqual(stage.groups[1].name, "group2") + + def test_inline_json_minimal_stages(self): + """Test inline JSON with minimal required fields.""" + minimal_data = { + "stages": [ + { + "name": "minimal-stage", + "groups": [ + {"name": "minimal-group"} + ] + } + ] + } + + inline_json = json.dumps(minimal_data) + result = get_update_run_strategy(self.mock_cmd, "fleet_update_runs", inline_json) + + # Verify minimal structure works + self.assertIsNotNone(result) + self.assertEqual(len(result.stages), 1) + + stage = result.stages[0] + self.assertEqual(stage.name, "minimal-stage") + self.assertEqual(stage.after_stage_wait_in_seconds, 0) # Should default to 0 + self.assertEqual(len(stage.groups), 1) + self.assertEqual(stage.groups[0].name, "minimal-group") + + def test_complex_stages_structure(self): + """Test more complex stages structure with multiple stages and groups.""" + complex_data = { + "stages": [ + { + "name": "stage1", + "groups": [ + {"name": "group1"}, + {"name": "group2"} + ], + "afterStageWaitInSeconds": 1800 + }, + { + "name": "stage2", + "groups": [ + {"name": "group3"} + ], + "afterStageWaitInSeconds": 3600 + } + ] + } + + inline_json = json.dumps(complex_data) + result = get_update_run_strategy(self.mock_cmd, "fleet_update_runs", inline_json) + + # Verify complex structure + self.assertIsNotNone(result) + self.assertEqual(len(result.stages), 2) + + # Verify first stage + stage1 = result.stages[0] + self.assertEqual(stage1.name, "stage1") + self.assertEqual(stage1.after_stage_wait_in_seconds, 1800) + self.assertEqual(len(stage1.groups), 2) + self.assertEqual(stage1.groups[0].name, "group1") + self.assertEqual(stage1.groups[1].name, "group2") + + # Verify second stage + stage2 = result.stages[1] + self.assertEqual(stage2.name, "stage2") + self.assertEqual(stage2.after_stage_wait_in_seconds, 3600) + self.assertEqual(len(stage2.groups), 1) + self.assertEqual(stage2.groups[0].name, "group3") + + def test_none_stages_returns_none(self): + """Test that None stages input returns None.""" + result = get_update_run_strategy(self.mock_cmd, "fleet_update_runs", None) + self.assertIsNone(result) + + def test_invalid_json_raises_error(self): + """Test that invalid JSON strings raise appropriate errors.""" + invalid_json = '{"stages": [{"name": "test", invalid_syntax}]}' + + # Should raise an error when parsing invalid JSON + with self.assertRaises(InvalidArgumentValueError): + get_update_run_strategy(self.mock_cmd, "fleet_update_runs", invalid_json) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/fleet/setup.py b/src/fleet/setup.py index 4f3c36256d4..c8f78141e8b 100644 --- a/src/fleet/setup.py +++ b/src/fleet/setup.py @@ -16,7 +16,7 @@ # TODO: Confirm this is the right version number you want and it matches your # HISTORY.rst entry. -VERSION = '1.6.0' +VERSION = '1.6.1' # The full list of classifiers is available at # https://pypi.python.org/pypi?%3Aaction=list_classifiers