Fix issues during enabling ipv6

Author: LiliDeng

lisa/sut_orchestrator/azure/arm_template.bicep

@@ -273,11 +273,11 @@ resource nodes_public_ip 'Microsoft.Network/publicIPAddresses@2020-05-01' = [for
   tags: tags
   name: '${nodes[i].name}-public-ip'
   properties: {
-    publicIPAllocationMethod: ((is_ultradisk || use_availability_zones) ? 'Static' : 'Dynamic')
+    publicIPAllocationMethod: ((is_ultradisk || use_availability_zones || use_ipv6) ? 'Static' : 'Dynamic')
     ipTags: (empty(ip_tags) ? null : ip_tags)
   }
   sku: {
-    name: ((is_ultradisk || use_availability_zones) ? 'Standard' : 'Basic')
+    name: ((is_ultradisk || use_availability_zones || use_ipv6) ? 'Standard' : 'Basic')
   }
   zones: (use_availability_zones ? availability_zones : null)
 }]

lisa/sut_orchestrator/azure/autogen_arm_template.json

@@ -6,7 +6,7 @@
     "_generator": {
       "name": "bicep",
       "version": "0.34.1.11899",
-      "templateHash": "13761993183944386189"
+      "templateHash": "6209161677269277510"
     }
   },
   "functions": [
@@ -607,11 +607,11 @@
       "location": "[parameters('location')]",
       "tags": "[parameters('tags')]",
       "properties": {
-        "publicIPAllocationMethod": "[if(or(parameters('is_ultradisk'), variables('use_availability_zones')), 'Static', 'Dynamic')]",
+        "publicIPAllocationMethod": "[if(or(or(parameters('is_ultradisk'), variables('use_availability_zones')), parameters('use_ipv6')), 'Static', 'Dynamic')]",
         "ipTags": "[if(empty(variables('ip_tags')), null(), variables('ip_tags'))]"
       },
       "sku": {
-        "name": "[if(or(parameters('is_ultradisk'), variables('use_availability_zones')), 'Standard', 'Basic')]"
+        "name": "[if(or(or(parameters('is_ultradisk'), variables('use_availability_zones')), parameters('use_ipv6')), 'Standard', 'Basic')]"
       },
       "zones": "[if(variables('use_availability_zones'), variables('availability_zones'), null())]"
     },

lisa/util/shell.py

@@ -1,7 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.
 
-import ipaddress
 import logging
 import os
 import re
@@ -78,9 +77,7 @@ def wait_tcp_port_ready(
         # If it's True, it means the direct connection doesn't work. Return a
         # mock value for test purpose.
         return True, 0
-    address_family = socket.AF_INET
-    if _is_valid_ipv6(address):
-        address_family = socket.AF_INET6
+    address_family = _get_address_family(address)
     timeout_timer = create_timer()
     while timeout_timer.elapsed(False) < timeout:
         with socket.socket(address_family, socket.SOCK_STREAM) as tcp_socket:
@@ -104,12 +101,20 @@ def wait_tcp_port_ready(
     return is_ready, result
 
 
-def _is_valid_ipv6(address: str) -> bool:
+def _get_address_family(address: str) -> Any:
     try:
-        ipaddress.IPv4Address(address)
-        return False
-    except ipaddress.AddressValueError:
-        return True
+        addr_info = socket.getaddrinfo(address, None)
+        for info in addr_info:
+            family = info[0]
+            if family == socket.AF_INET:
+                return socket.AF_INET
+            elif family == socket.AF_INET6:
+                return socket.AF_INET6
+        return socket.AF_INET
+    except socket.gaierror:
+        return socket.AF_INET
+    except Exception:
+        return socket.AF_INET
 
 
 class WindowsShellType(object):