Add flag for testing windows tests over the internal IP.

PiperOrigin-RevId: 315987647

CHANGES.next.md

@@ -65,7 +65,7 @@
 -   Update Boto client library for S3 to boto3.
 -   Add MultiStreamDelete to Object Storage API Tests.
 -   Add GCE Placement Group support. By default, placement groups are not
-      created.
+    created.
 -   Add ICMP ping support for Azure using EXTERNAL ip address.
 
 ### Enhancements:
@@ -110,6 +110,9 @@
 -   Add Spec17 configs compatible with v1.1.
 -   Better support for pytyping.
     -   Make BaseVirtualMachine inherit from BaseOsMixin.
+-   Add connect_via_internal_ip flag for testing windows tests over the internal
+    IP, similar to the ssh_via_internal_ip flag for linux.
+    -   Deprecate `--ssh_via_internal_ip` in favor of this new flag.
 
 ### Bug fixes and maintenance updates:
 

perfkitbenchmarker/linux_virtual_machine.py

@@ -133,13 +133,6 @@
 flags.DEFINE_integer(
     'scp_connect_timeout', 30, 'timeout for SCP connection.', lower_bound=0)
 
-flags.DEFINE_boolean(
-    'ssh_via_internal_ip', False,
-    'Whether to use internal IP addresses for running commands on and pushing '
-    'data to VMs. By default, PKB interacts with VMs using external IP '
-    'addresses.'
-)
-
 flags.DEFINE_string(
     'append_kernel_command_line', None,
     'String to append to the kernel command line. The presence of any '
@@ -755,8 +748,7 @@ def RemoteHostCopy(self, file_path, remote_path='', copy_to=True):
         file_path = file_path.split(':', 1)[1]
       # Replace the last instance of '\' with '/' to make scp happy.
       file_path = '/'.join(file_path.rsplit('\\', 1))
-    remote_ip = '[%s]' % (
-        self.internal_ip if FLAGS.ssh_via_internal_ip else self.ip_address)
+    remote_ip = '[%s]' % self.GetConnectionIp()
     remote_location = '%s@%s:%s' % (
         self.user_name, remote_ip, remote_path)
     scp_cmd = ['scp', '-P', str(self.ssh_port), '-pr']
@@ -851,8 +843,7 @@ def RemoteHostCommandWithReturnCode(self,
       # Multi-line commands passed to ssh won't work on Windows unless the
       # newlines are escaped.
       command = command.replace('\n', '\\n')
-    ip_address = (
-        self.internal_ip if FLAGS.ssh_via_internal_ip else self.ip_address)
+    ip_address = self.GetConnectionIp()
     user_host = '%s@%s' % (self.user_name, ip_address)
     ssh_cmd = ['ssh', '-A', '-p', str(self.ssh_port), user_host]
     ssh_private_key = (self.ssh_private_key if self.is_static else

perfkitbenchmarker/virtual_machine.py

@@ -109,6 +109,18 @@ def ValidateVmMetadataFlag(options_list):
     'preprovision_ignore_checksum', False,
     'Ignore checksum verification for preprovisioned data. '
     'Not recommended, please use with caution')
+flags.DEFINE_boolean(
+    'connect_via_internal_ip', False,
+    'Whether to use internal IP addresses for running commands on and pushing '
+    'data to VMs. By default, PKB interacts with VMs using external IP '
+    'addresses.')
+
+# Deprecated. Use connect_via_internal_ip.
+flags.DEFINE_boolean(
+    'ssh_via_internal_ip', False,
+    'Whether to use internal IP addresses for running commands on and pushing '
+    'data to VMs. By default, PKB interacts with VMs using external IP '
+    'addresses.')
 
 # Note: If adding a gpu type here, be sure to add it to
 # the flag definition in pkb.py too.
@@ -880,6 +892,12 @@ def __str__(self):
       return self.ip_address
     return super(BaseVirtualMachine, self).__str__()
 
+  def GetConnectionIp(self):
+    """Gets the IP to use for connecting to the VM."""
+    if FLAGS.ssh_via_internal_ip or FLAGS.connect_via_internal_ip:
+      return self.internal_ip
+    return self.ip_address
+
   def CreateScratchDisk(self, disk_spec):
     """Create a VM's scratch disk.
 

perfkitbenchmarker/windows_virtual_machine.py

@@ -205,9 +205,10 @@ def RemoteCommand(self, command, should_log=False, ignore_failure=False,
           command timed out.
     """
     logging.info('Running command on %s: %s', self, command)
-    s = winrm.Session('https://%s:%s' % (self.ip_address, self.winrm_port),
-                      auth=(self.user_name, self.password),
-                      server_cert_validation='ignore')
+    s = winrm.Session(
+        'https://%s:%s' % (self.GetConnectionIp(), self.winrm_port),
+        auth=(self.user_name, self.password),
+        server_cert_validation='ignore')
     encoded_command = six.ensure_str(
         base64.b64encode(command.encode('utf_16_le')))
 
@@ -309,7 +310,7 @@ def RemoteCopy(self, local_path, remote_path='', copy_to=True):
 
     drive, remote_path = ntpath.splitdrive(remote_path)
     remote_drive = (drive or self.system_drive).rstrip(':')
-    network_drive = '\\\\%s\\%s$' % (self.ip_address, remote_drive)
+    network_drive = '\\\\%s\\%s$' % (self.GetConnectionIp(), remote_drive)
 
     if vm_util.RunningOnWindows():
       self._PsDriveRemoteCopy(local_path, remote_path, copy_to, network_drive)