Fix unit test issue where az vm run-command isn't running anymore (#517)

* Fix vm run-command syntax for ci-k8s-common.sh

Signed-off-by: Mark Rossett <marosset@microsoft.com>

* updates

* azure-cli 2.76

---------

Signed-off-by: Mark Rossett <marosset@microsoft.com>

Author: marosset

scripts/ci-k8s-common.sh

@@ -14,26 +14,36 @@ function onError(){
 
 ensure_azure_cli() {
     if [[ -z "$(command -v az)" ]]; then
-        echo "installing Azure CLI"
+        echo "installing Azure CLI v2.76.0"
         apt-get update && apt-get install -y ca-certificates curl apt-transport-https lsb-release gnupg
         curl -sL https://packages.microsoft.com/keys/microsoft.asc | gpg --dearmor | tee /etc/apt/trusted.gpg.d/microsoft.gpg > /dev/null
         AZ_REPO=$(lsb_release -cs)
         echo "deb [arch=amd64] https://packages.microsoft.com/repos/azure-cli/ ${AZ_REPO} main" | tee /etc/apt/sources.list.d/azure-cli.list
-        apt-get update && apt-get install -y azure-cli
-  	
-        if [[ -n "${AZURE_FEDERATED_TOKEN_FILE:-}" ]]; then
-            echo "Logging in with federated token"
-            # AZURE_CLIENT_ID has been overloaded with Azure Workload ID in the preset-azure-cred-wi.
-            # This is done to avoid exporting Azure Workload ID as AZURE_CLIENT_ID in the test scenarios.
-            az login --service-principal -u "${AZURE_CLIENT_ID}" -t "${AZURE_TENANT_ID}" --federated-token "$(cat "${AZURE_FEDERATED_TOKEN_FILE}")" > /dev/null
-
-            # Use --auth-mode "login" in az storage commands to use RBAC permissions of login identity. This is a well known ENV variable the Azure cli
-            export AZURE_STORAGE_AUTH_MODE="login"
+        apt-get update && apt-get install -y azure-cli=2.76.0-1~${AZ_REPO}
+    else
+        # Check if we have the correct version
+        CURRENT_VERSION=$(az version --query '."azure-cli"' -o tsv 2>/dev/null || echo "unknown")
+        REQUIRED_VERSION="2.76.0"
+        if [[ "$CURRENT_VERSION" != "$REQUIRED_VERSION" ]]; then
+            echo "Warning: Azure CLI version is $CURRENT_VERSION, but $REQUIRED_VERSION is required"
+            echo "Consider running: apt-get install -y azure-cli=${REQUIRED_VERSION}-1~$(lsb_release -cs)"
         else
-            echo "AZURE_FEDERATED_TOKEN_FILE environment variable must be set to path location of token file"
-            exit 1
+            echo "Azure CLI version $CURRENT_VERSION is correct"
         fi
     fi
+    
+    if [[ -n "${AZURE_FEDERATED_TOKEN_FILE:-}" ]]; then
+        echo "Logging in with federated token"
+        # AZURE_CLIENT_ID has been overloaded with Azure Workload ID in the preset-azure-cred-wi.
+        # This is done to avoid exporting Azure Workload ID as AZURE_CLIENT_ID in the test scenarios.
+        az login --service-principal -u "${AZURE_CLIENT_ID}" -t "${AZURE_TENANT_ID}" --federated-token "$(cat "${AZURE_FEDERATED_TOKEN_FILE}")" > /dev/null
+
+        # Use --auth-mode "login" in az storage commands to use RBAC permissions of login identity. This is a well known ENV variable the Azure cli
+        export AZURE_STORAGE_AUTH_MODE="login"
+    else
+        echo "AZURE_FEDERATED_TOKEN_FILE environment variable must be set to path location of token file"
+        exit 1
+    fi
 }
 
 
@@ -148,11 +158,48 @@ run_remote_cmd() {
 
 enable_ssh_windows() {
     echo "Enabling SSH for Windows VM"
-    az vm run-command invoke  --command-id RunPowerShellScript -n ${VM_NAME} -g ${AZURE_RESOURCE_GROUP} --scripts "@$(pwd)/scripts/enable_ssh_windows.ps1" --parameters "SSHPublicKey=${AZURE_SSH_PUBLIC_KEY}"
+    local SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+    local ENABLE_SSH_WINDOWS_SCRIPT="${SCRIPT_DIR}/enable_ssh_windows.ps1"
+    echo "Using run-command script: ${ENABLE_SSH_WINDOWS_SCRIPT}"
+    if [ ! -f "${ENABLE_SSH_WINDOWS_SCRIPT}" ]; then
+        echo "Enable-SSH script not found at ${ENABLE_SSH_WINDOWS_SCRIPT}"
+        return 1
+    fi
+    local run_command_output
+    if ! run_command_output=$(az vm run-command invoke --command-id RunPowerShellScript \
+        -n ${VM_NAME} -g ${AZURE_RESOURCE_GROUP} \
+        --scripts @${ENABLE_SSH_WINDOWS_SCRIPT} \
+        --parameters "SSHPublicKey=${AZURE_SSH_PUBLIC_KEY}" \
+        --only-show-errors -o json 2>&1); then
+        echo "Failed to enable SSH on Windows VM"
+        echo "Azure CLI output:"
+        echo "${run_command_output}"
+        return 1
+    fi
+    echo "Raw Azure run-command output:"
+    printf '%s\n' "${run_command_output}"
+    echo "Azure run-command output:"
+    printf '%s\n' "${run_command_output}" | jq -r '.value[].message'
 }
 
 test_ssh_connection() {
-	echo "Testing ssh connection to Windows VM"
+    echo "Checking sshd service state on Windows VM"
+    local service_check_output
+    if ! service_check_output=$(az vm run-command invoke --command-id RunPowerShellScript \
+        -n ${VM_NAME} -g ${AZURE_RESOURCE_GROUP} \
+        --scripts 'param([string]$serviceName) $svc = Get-Service -Name $serviceName -ErrorAction Stop; Write-Output ("sshd service status: {0}" -f $svc.Status); if ($svc.Status -ne "Running") { throw "Service $serviceName is not running" }' \
+        --parameters "serviceName=sshd" \
+        --only-show-errors -o json 2>&1); then
+        echo "Azure run-command indicates sshd service is not running"
+        echo "Azure CLI output:"
+        echo "${service_check_output}"
+        exit 1
+    fi
+    echo "Raw Azure run-command output:" 
+    printf '%s\n' "${service_check_output}"
+    echo "Azure run-command output:"
+    printf '%s\n' "${service_check_output}" | jq -r '.value[].message'
+    echo "Testing ssh connection to Windows VM"
     SSH_KEY_FILE=.sshkey
 	if ! ssh -i ${SSH_KEY_FILE} ${SSH_OPTS} azureuser@${VM_PUB_IP}  "hostname";
     then