diff --git a/src/connectedk8s/azext_connectedk8s/custom.py b/src/connectedk8s/azext_connectedk8s/custom.py new file mode 100644 index 00000000000..bdd2961e2c3 --- /dev/null +++ b/src/connectedk8s/azext_connectedk8s/custom.py @@ -0,0 +1,101 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +""" +Custom commands for connectedk8s extension +""" + +import logging +from knack.util import CLIError + +logger = logging.getLogger(__name__) + + +def enable_features(cmd, client, resource_group_name, cluster_name, features=None, + azure_rbac_enabled=False, custom_locations_enabled=False): + """ + Enable features for connected Kubernetes cluster + + :param cmd: Command context + :param client: Connected cluster client + :param resource_group_name: Resource group name + :param cluster_name: Cluster name + :param features: List of features to enable + :param azure_rbac_enabled: Enable Azure RBAC + :param custom_locations_enabled: Enable custom locations + """ + + # Parse features + enable_azure_rbac = azure_rbac_enabled + enable_cl = custom_locations_enabled + + if features: + if 'azure-rbac' in features: + enable_azure_rbac = True + if 'custom-locations' in features: + enable_cl = True + + logger.info(f"Enabling features: azure-rbac={enable_azure_rbac}, custom-locations={enable_cl}") + + # Initialize variables with safe defaults before conditional assignment + final_enable_cl = False + custom_locations_oid = "" + + # Other setup code... + helm_upgrade_command = ["helm", "upgrade", "azure-arc", "azure-arc"] + + # Enable Azure RBAC if requested + if enable_azure_rbac: + helm_upgrade_command.extend(["--set", "systemDefaultValues.guard.enabled=true"]) + logger.info("Azure RBAC feature enabled") + + # Enable custom locations if requested + if enable_cl: + # This is where the variables are defined - only in this block + final_enable_cl = True + custom_locations_oid = "some-object-id-value" + + helm_upgrade_command.extend([ + "--set", f"systemDefaultValues.customLocations.enabled={final_enable_cl}", + "--set", f"systemDefaultValues.customLocations.oid={custom_locations_oid}" + ]) + logger.info("Custom locations feature enabled") + + # Additional helm upgrade command setup... + helm_upgrade_command.extend(["--namespace", "azure-arc"]) + + # BUG: These variables are referenced here but only defined inside the if enable_cl block above + # This will cause UnboundLocalError when enable_cl is False (i.e., when only azure-rbac is enabled) + final_command = helm_upgrade_command + [ + "--set", f"systemDefaultValues.customLocations.finalEnabled={final_enable_cl}", + "--set", f"systemDefaultValues.customLocations.finalOid={custom_locations_oid}" + ] + + logger.info(f"Running helm command: {' '.join(final_command)}") + + # Execute the helm command (simulated) + try: + # subprocess.run(final_command, check=True) + logger.info("Helm upgrade completed successfully") + return {"status": "success", "message": "Features enabled successfully"} + except Exception as e: + raise CLIError(f"Failed to enable features: {str(e)}") + + +def list_features(cmd, client, resource_group_name, cluster_name): + """ + List enabled features for connected Kubernetes cluster + """ + # Implementation for listing features + return {"azure-rbac": False, "custom-locations": False} + + +def disable_features(cmd, client, resource_group_name, cluster_name, features=None): + """ + Disable features for connected Kubernetes cluster + """ + # Implementation for disabling features + logger.info(f"Disabling features: {features}") + return {"status": "success", "message": "Features disabled successfully"} \ No newline at end of file diff --git a/src/connectedk8s/azext_connectedk8s/tests/__init__.py b/src/connectedk8s/azext_connectedk8s/tests/__init__.py new file mode 100644 index 00000000000..6782f1bf674 --- /dev/null +++ b/src/connectedk8s/azext_connectedk8s/tests/__init__.py @@ -0,0 +1 @@ +# Tests for connectedk8s extension \ No newline at end of file diff --git a/src/connectedk8s/azext_connectedk8s/tests/test_custom.py b/src/connectedk8s/azext_connectedk8s/tests/test_custom.py new file mode 100644 index 00000000000..139ba99ba9d --- /dev/null +++ b/src/connectedk8s/azext_connectedk8s/tests/test_custom.py @@ -0,0 +1,103 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import unittest +from unittest.mock import Mock, MagicMock + + +class TestEnableFeatures(unittest.TestCase): + """ + Test cases for the enable_features function to ensure UnboundLocalError is fixed + """ + + def setUp(self): + """Set up test fixtures""" + self.mock_cmd = Mock() + self.mock_client = Mock() + self.resource_group_name = "test-rg" + self.cluster_name = "test-cluster" + + def test_enable_azure_rbac_only_no_unbound_error(self): + """ + Test that enabling only azure-rbac does not cause UnboundLocalError + This is a regression test for the bug fix. + """ + from azext_connectedk8s.custom import enable_features + + # This should not raise UnboundLocalError after the fix + try: + result = enable_features( + cmd=self.mock_cmd, + client=self.mock_client, + resource_group_name=self.resource_group_name, + cluster_name=self.cluster_name, + features=["azure-rbac"] + ) + # If we reach here, the function completed without UnboundLocalError + self.assertIsNotNone(result) + self.assertEqual(result["status"], "success") + except UnboundLocalError as e: + self.fail(f"UnboundLocalError should be fixed but was raised: {e}") + + def test_enable_custom_locations_only(self): + """ + Test enabling only custom-locations works correctly + """ + from azext_connectedk8s.custom import enable_features + + try: + result = enable_features( + cmd=self.mock_cmd, + client=self.mock_client, + resource_group_name=self.resource_group_name, + cluster_name=self.cluster_name, + features=["custom-locations"] + ) + self.assertIsNotNone(result) + self.assertEqual(result["status"], "success") + except Exception as e: + self.fail(f"Unexpected error when enabling custom-locations only: {e}") + + def test_enable_both_features(self): + """ + Test enabling both azure-rbac and custom-locations works correctly + """ + from azext_connectedk8s.custom import enable_features + + try: + result = enable_features( + cmd=self.mock_cmd, + client=self.mock_client, + resource_group_name=self.resource_group_name, + cluster_name=self.cluster_name, + features=["azure-rbac", "custom-locations"] + ) + self.assertIsNotNone(result) + self.assertEqual(result["status"], "success") + except Exception as e: + self.fail(f"Unexpected error when enabling both features: {e}") + + def test_enable_no_features(self): + """ + Test calling enable_features without any features specified + """ + from azext_connectedk8s.custom import enable_features + + try: + result = enable_features( + cmd=self.mock_cmd, + client=self.mock_client, + resource_group_name=self.resource_group_name, + cluster_name=self.cluster_name, + features=[] + ) + self.assertIsNotNone(result) + self.assertEqual(result["status"], "success") + except Exception as e: + self.fail(f"Unexpected error when enabling no features: {e}") + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file