From 04ee09851bf4726887fa3ffac379dbb2a2343b8c Mon Sep 17 00:00:00 2001 From: Jessie Wang Date: Thu, 13 Aug 2026 10:59:20 +1000 Subject: [PATCH] Add namespace access for list and rotate keys --- src/aimanager/azext_aimanager/_help.py | 19 +++++++ src/aimanager/azext_aimanager/_params.py | 6 ++ src/aimanager/azext_aimanager/commands.py | 2 + src/aimanager/azext_aimanager/custom.py | 14 +++++ .../test_namespace_accesskeys_scenario.py | 56 +++++++++++++++++++ 5 files changed, 97 insertions(+) create mode 100644 src/aimanager/azext_aimanager/tests/latest/test_namespace_accesskeys_scenario.py diff --git a/src/aimanager/azext_aimanager/_help.py b/src/aimanager/azext_aimanager/_help.py index 4964224cc91..c98431ee4a0 100644 --- a/src/aimanager/azext_aimanager/_help.py +++ b/src/aimanager/azext_aimanager/_help.py @@ -141,6 +141,25 @@ text: az aimanager namespace get-credentials -m my-ai-manager -g myrg --name team-alpha -f - """ +helps['aimanager namespace list-accesskeys'] = """ + type: command + short-summary: List the inference gateway endpoint and API keys for an AI Manager namespace. + examples: + - name: List namespace access keys + text: az aimanager namespace list-accesskeys -m my-ai-manager -g myrg --name team-alpha +""" + +helps['aimanager namespace rotate-accesskeys'] = """ + type: command + short-summary: Rotate the API keys for an AI Manager namespace. + long-summary: |- + Generates a new primary key and moves the previous primary key to the secondary key so + clients can roll over without downtime. + examples: + - name: Rotate namespace access keys + text: az aimanager namespace rotate-accesskeys -m my-ai-manager -g myrg --name team-alpha +""" + helps['aimanager namespace modeldeployment'] = """ type: group short-summary: Manage model deployments within an AI Manager namespace. diff --git a/src/aimanager/azext_aimanager/_params.py b/src/aimanager/azext_aimanager/_params.py index b859b8577de..993166d9a8c 100644 --- a/src/aimanager/azext_aimanager/_params.py +++ b/src/aimanager/azext_aimanager/_params.py @@ -80,6 +80,12 @@ def load_arguments(self, _): c.argument('aks_custom_headers', options_list=['--aks-custom-headers'], help='Comma-separated key=value pairs to specify custom headers.') + for scope in ['aimanager namespace list-accesskeys', + 'aimanager namespace rotate-accesskeys']: + with self.argument_context(scope) as c: + c.argument('aks_custom_headers', options_list=['--aks-custom-headers'], + help='Comma-separated key=value pairs to specify custom headers.') + with self.argument_context('aimanager namespace modeldeployment') as c: c.argument('ai_manager_name', options_list=['--aimanager-name'], validator=validate_ai_manager_name, diff --git a/src/aimanager/azext_aimanager/commands.py b/src/aimanager/azext_aimanager/commands.py index d100c263b27..c0331e54dc2 100644 --- a/src/aimanager/azext_aimanager/commands.py +++ b/src/aimanager/azext_aimanager/commands.py @@ -50,6 +50,8 @@ def load_command_table(self, _): g.custom_command("list", "list_aimanager_namespace") g.custom_command("delete", "delete_aimanager_namespace", supports_no_wait=True, confirmation=True) g.custom_command("get-credentials", "aimanager_namespace_get_credentials") + g.custom_command("list-accesskeys", "list_aimanager_namespace_access_keys") + g.custom_command("rotate-accesskeys", "rotate_aimanager_namespace_access_keys", confirmation=True) g.wait_command("wait") # aimanager namespace modeldeployment command group diff --git a/src/aimanager/azext_aimanager/custom.py b/src/aimanager/azext_aimanager/custom.py index 6da27d25606..349696c708f 100644 --- a/src/aimanager/azext_aimanager/custom.py +++ b/src/aimanager/azext_aimanager/custom.py @@ -299,6 +299,20 @@ def aimanager_namespace_get_credentials(cmd, resource_group_name, ai_manager_name, namespace_name, headers=headers) _write_kubeconfig(credential_results, path, overwrite_existing, context_name) + +def list_aimanager_namespace_access_keys(cmd, client, resource_group_name, ai_manager_name, + namespace_name, aks_custom_headers=None): # pylint: disable=unused-argument + headers = get_aks_custom_headers(aks_custom_headers) + return client.list_access_keys( + resource_group_name, ai_manager_name, namespace_name, headers=headers) + + +def rotate_aimanager_namespace_access_keys(cmd, client, resource_group_name, ai_manager_name, + namespace_name, aks_custom_headers=None): # pylint: disable=unused-argument + headers = get_aks_custom_headers(aks_custom_headers) + return client.rotate_keys( + resource_group_name, ai_manager_name, namespace_name, headers=headers) + # endregion diff --git a/src/aimanager/azext_aimanager/tests/latest/test_namespace_accesskeys_scenario.py b/src/aimanager/azext_aimanager/tests/latest/test_namespace_accesskeys_scenario.py new file mode 100644 index 00000000000..9275c73a872 --- /dev/null +++ b/src/aimanager/azext_aimanager/tests/latest/test_namespace_accesskeys_scenario.py @@ -0,0 +1,56 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from unittest.mock import MagicMock, patch + +from azure.cli.testsdk import ScenarioTest + +from azext_aimanager.vendored_sdks.v2026_05_02_preview import models + + +class NamespaceAccessKeysScenarioTest(ScenarioTest): + + def test_namespace_access_key_commands(self): + current_access = models.NamespaceAccessInfo({ + 'endpoint': 'https://namespace.example/v1', + 'primaryKey': 'primary-before', + 'secondaryKey': 'secondary-before', + }) + rotated_access = models.NamespaceAccessInfo({ + 'endpoint': 'https://namespace.example/v1', + 'primaryKey': 'primary-after', + 'secondaryKey': 'primary-before', + 'lastRotatedAt': '2026-08-12T00:00:00Z', + }) + + operations = MagicMock() + operations.list_access_keys.return_value = current_access + operations.rotate_keys.return_value = rotated_access + service_client = MagicMock() + service_client.ai_manager_namespaces = operations + + command_suffix = '-g rg -m manager -n namespace --aks-custom-headers test-header=value' + with patch('azext_aimanager._client_factory.get_aimanager_client', + return_value=service_client): + self.cmd( + 'aimanager namespace list-accesskeys ' + command_suffix, + checks=[ + self.check('endpoint', 'https://namespace.example/v1'), + self.check('primaryKey', 'primary-before'), + self.check('secondaryKey', 'secondary-before'), + ]) + self.cmd( + 'aimanager namespace rotate-accesskeys ' + command_suffix + ' --yes', + checks=[ + self.check('primaryKey', 'primary-after'), + self.check('secondaryKey', 'primary-before'), + self.check('lastRotatedAt', '2026-08-12T00:00:00+00:00'), + ]) + + expected_headers = {'test-header': 'value'} + operations.list_access_keys.assert_called_once_with( + 'rg', 'manager', 'namespace', headers=expected_headers) + operations.rotate_keys.assert_called_once_with( + 'rg', 'manager', 'namespace', headers=expected_headers)