ctest: add ssh_port option

Add --ssh_port option to ctest to specify which port should
we used to access the image.

This option helps us to launch different VM used for VMTest so that we
can running them in parallel.

BUG=None
TEST=tryjob betty-vmtest-informational
TEST=manual run ctest
TEST=unittest

Change-Id: Id183da11f6eb9e2b5bc69477055b0119185552d3
Reviewed-on: https://chromium-review.googlesource.com/724276
Commit-Ready: Po-Hsien Wang <pwang@chromium.org>
Tested-by: Po-Hsien Wang <pwang@chromium.org>
Reviewed-by: Ilja H. Friedel <ihf@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/au_test_harness/au_test.py b/au_test_harness/au_test.py
index d78bff2..97300a0 100644
--- a/au_test_harness/au_test.py
+++ b/au_test_harness/au_test.py
@@ -131,7 +131,7 @@
     This test checks that we can update by updating the stateful partition
     rather than wiping it.
     """
-    self.worker.Initialize(9222)
+    self.worker.Initialize(self.options.ssh_port or 9222)
     # Just make sure some tests pass on original image.  Some old images
     # don't pass many tests.
     base_image_path = self.worker.PrepareBase(self.base_image_path)
@@ -150,7 +150,7 @@
     This test checks that we can update successfully after wiping the
     stateful partition.
     """
-    self.worker.Initialize(9223)
+    self.worker.Initialize(self.options.ssh_port or 9223)
     # Just make sure some tests pass on original image.  Some old images
     # don't pass many tests.
     base_image_path = self.worker.PrepareBase(self.base_image_path)
@@ -195,12 +195,12 @@
         self.data_size += len(data)
         return data
 
-    self.worker.Initialize(9224)
+    self.worker.Initialize(self.options.ssh_port or 9224)
     self.AttemptUpdateWithFilter(InterruptionFilter(), proxy_port=8082)
 
   def testSimpleSignedUpdate(self):
     """Test that updates to itself with a signed payload."""
-    self.worker.Initialize(9226)
+    self.worker.Initialize(self.options.ssh_port or 9226)
     signed_target_image_path = self.worker.PrepareBase(self.target_image_path,
                                                        signed_base=True)
     if self.payload_signing_key:
@@ -216,7 +216,7 @@
     We explicitly don't use test prefix so that isn't run by default.  Can be
     run using test_prefix option.
     """
-    self.worker.Initialize(9227)
+    self.worker.Initialize(self.options.ssh_port or 9227)
     target_image_path = self.worker.PrepareBase(self.target_image_path)
     self.worker.PerformUpdate(target_image_path, target_image_path)
     self.assertTrue(self.worker.VerifyImage(self))
@@ -227,7 +227,7 @@
     We explicitly don't use test prefix so that isn't run by default.  Can be
     run using test_prefix option.
     """
-    self.worker.Initialize(9228)
+    self.worker.Initialize(self.options.ssh_port or 9228)
     self.worker.PrepareBase(self.target_image_path)
     self.assertTrue(self.worker.VerifyImage(self))
 
@@ -263,19 +263,19 @@
         self.data_size += len(data)
         return data
 
-    self.worker.Initialize(9225)
+    self.worker.Initialize(self.options.ssh_port or 9225)
     self.AttemptUpdateWithFilter(DelayedFilter(), proxy_port=8083)
 
   def NotestPlatformToolchainOptions(self):
     """Tests the hardened toolchain options."""
-    self.worker.Initialize(9229)
+    self.worker.Initialize(self.options.ssh_port or 9229)
     self.worker.PrepareBase(self.base_image_path)
     self.assertTrue(self.worker.VerifyImage('platform_ToolchainOptions'))
 
   # TODO(sosa): Get test to work with verbose.
   def NotestPartialUpdate(self):
     """Tests what happens if we attempt to update with a truncated payload."""
-    self.worker.Initialize(9230)
+    self.worker.Initialize(self.options.ssh_port or 9230)
     # Preload with the version we are trying to test.
     self.worker.PrepareBase(self.target_image_path)
 
@@ -294,7 +294,7 @@
   # TODO(sosa): Get test to work with verbose.
   def NotestCorruptedUpdate(self):
     """Tests what happens if we attempt to update with a corrupted payload."""
-    self.worker.Initialize(9231)
+    self.worker.Initialize(self.options.ssh_port or 9231)
     # Preload with the version we are trying to test.
     self.worker.PrepareBase(self.target_image_path)
 
diff --git a/au_test_harness/au_worker.py b/au_test_harness/au_worker.py
index 4b3417a..8bc16d0 100644
--- a/au_test_harness/au_worker.py
+++ b/au_test_harness/au_worker.py
@@ -272,7 +272,7 @@
     """
     # Initialize port vars.
     self._ssh_port = port
-    self._kvm_pid_file = '/tmp/kvm.%d' % port
+    self._kvm_pid_file = '/tmp/kvm.%d' % self._ssh_port
 
     # Initialize test results directory.
     self.test_name = inspect.stack()[1][3]
diff --git a/au_test_harness/cros_au_test_harness.py b/au_test_harness/cros_au_test_harness.py
index cd5f3c9..e7c3ecb 100755
--- a/au_test_harness/cros_au_test_harness.py
+++ b/au_test_harness/cros_au_test_harness.py
@@ -153,6 +153,12 @@
   if opts.ssh_private_key and not os.path.isfile(opts.ssh_private_key):
     parser.error('Testing requires a valid path to the ssh private key.')
 
+  if opts.ssh_port and opts.ssh_port < 1024:
+    parser.error('Testing requires a valid port higher than 1024.')
+
+  if opts.ssh_port and not opts.test_prefix:
+    parser.error('Testing with ssh_port requires test_prefix specified.')
+
   if opts.test_results_root:
     if not 'chroot/tmp' in opts.test_results_root:
       parser.error('Must specify a test results root inside tmp in a chroot.')
@@ -213,6 +219,9 @@
   parser.add_argument('--ssh_private_key', default=None,
                       help='Path to the private key to use to ssh into the '
                       'image as the root user.')
+  parser.add_argument('--ssh_port', default=None, type=int,
+                      help='ssh port used to ssh into image. (Should only be'
+                      ' used with --test_prefix)')
   opts = parser.parse_args()
 
   CheckOpts(parser, opts)
diff --git a/ctest/ctest.py b/ctest/ctest.py
index d14a3b1..9d5f2d1 100755
--- a/ctest/ctest.py
+++ b/ctest/ctest.py
@@ -81,6 +81,7 @@
 
     # An optional ssh private key used for testing.
     self.ssh_private_key = opts.ssh_private_key
+    self.ssh_port = opts.ssh_port
 
   def GeneratePublicKey(self):
     """Returns the path to a generated public key from the UE private key."""
@@ -185,6 +186,8 @@
 
     if self.ssh_private_key is not None:
       cmd.append('--ssh_private_key=%s' % self.ssh_private_key)
+    if self.ssh_port is not None:
+      cmd.append('--ssh_port=%s' % self.ssh_port)
 
     if suite:
       cmd.append('--verify_suite_name=%s' % suite)
@@ -254,6 +257,9 @@
   parser.add_argument('--ssh_private_key', default=None,
                       help='Path to the private key to use to ssh into the '
                       'image as the root user')
+  parser.add_argument('--ssh_port', default=None, type=int,
+                      help='ssh port used to ssh into image. (Should only be'
+                      ' used with either --quick_update or --only_verify)')
 
   opts = parser.parse_args()
 
@@ -262,6 +268,10 @@
   if opts.only_verify and opts.quick_update:
     parser.error(
         'Only one of --only_verify or --quick_update should be specified.')
+  if opts.ssh_port and not (opts.only_verify or opts.quick_update):
+    parser.error(
+        'ssh_port should be specified with either --only_verify or '
+        '--quick_update')
 
   # force absolute path for these opts, since a chdir occurs deeper in the
   # codebase.