Create a transparent proxy for test programs to use, with the

bin/cros_au_test_harness.py

Index: bin/cros_au_test_harness.py
diff --git a/bin/cros_au_test_harness.py b/bin/cros_au_test_harness.py
index eacf837d1856df1d01b555ca2ee8eb9e8d59c125..d7d655a87149ceec11b0df480667ba16ba4b1812 100755
--- a/bin/cros_au_test_harness.py
+++ b/bin/cros_au_test_harness.py
@@ -8,6 +8,8 @@ import optparse
 import os
 import re
 import sys
+import thread
+import time
 import unittest
 import urllib
 
@@ -19,6 +21,8 @@ from cros_build_lib import RunCommand
 from cros_build_lib import RunCommandCaptureOutput
 from cros_build_lib import Warning
 
+import cros_test_proxy
+
 # VM Constants.
 _FULL_VDISK_SIZE = 6072
 _FULL_STATEFULFS_SIZE = 3074
@@ -87,13 +91,14 @@ class AUTest(object):
     else:
       self._UpdateImageReportError(image)
 
-  def _UpdateImageReportError(self, image_path, stateful_change='old'):
+  def _UpdateImageReportError(self, image_path, stateful_change='old',
+                              proxy_port=None):
     """Calls UpdateImage and reports any error to the console.
     
        Still throws the exception.
     """
     try:
-      self.UpdateImage(image_path, stateful_change)
+      self.UpdateImage(image_path, stateful_change, proxy_port)
     except UpdateException as err:
       # If the update fails, print it out
       Warning(err.stdout)
@@ -116,7 +121,7 @@ class AUTest(object):
     """Prepares target with base_image_path."""
     pass
 
-  def UpdateImage(self, image_path, stateful_change='old'):
+  def UpdateImage(self, image_path, stateful_change='old', proxy_port=None):
     """Updates target with the image given by the image_path.
 
     Args:
@@ -129,7 +134,7 @@ class AUTest(object):
     """
     pass
 
-  def UpdateUsingPayload(self, update_path, stateful_change='old'):
+  def UpdateUsingPayload(self, update_path, stateful_change='old', proxy_port=None):

[adlr, 2010/12/07 02:26:51]

80 cols

[dgarrett, 2010/12/07 03:48:18]

Done.

     """Updates target with the pre-generated update stored in update_path
 
     Args:
@@ -183,7 +188,7 @@ class AUTest(object):
     """
     # Just make sure some tests pass on original image.  Some old images
     # don't pass many tests.
-    self.PrepareBase(image_path=base_image_path)
+    self.PrepareBase(base_image_path)
     # TODO(sosa): move to 100% once we start testing using the autotest paired
     # with the dev channel.
     percent_passed = self.VerifyImage(10)
@@ -206,7 +211,7 @@ class AUTest(object):
     """
     # Just make sure some tests pass on original image.  Some old images
     # don't pass many tests.
-    self.PrepareBase(image_path=base_image_path)
+    self.PrepareBase(base_image_path)
     # TODO(sosa): move to 100% once we start testing using the autotest paired
     # with the dev channel.
     percent_passed = self.VerifyImage(10)
@@ -224,7 +229,7 @@ class AUTest(object):
   def testPartialUpdate(self):
     """Tests what happens if we attempt to update with a truncated payload."""
     # Preload with the version we are trying to test.
-    self.PrepareBase(image_path=target_image_path)
+    self.PrepareBase(target_image_path)
 
     # Image can be updated at:
     # ~chrome-eng/chromeos/localmirror/autest-images
@@ -241,7 +246,7 @@ class AUTest(object):
   def testCorruptedUpdate(self):
     """Tests what happens if we attempt to update with a corrupted payload."""
     # Preload with the version we are trying to test.
-    self.PrepareBase(image_path=target_image_path)
+    self.PrepareBase(target_image_path)
 
     # Image can be updated at:
     # ~chrome-eng/chromeos/localmirror/autest-images
@@ -256,6 +261,92 @@ class AUTest(object):
     expected_msg='zlib inflate() error:-3'
     self._AttemptUpdateWithPayloadExpectedFailure(payload, expected_msg)
 
+  def testInterruptedUpdate(self):
+    """Tests what happens if we interrupt payload delivery 3 times."""
+
+    class InterruptionFilter(cros_test_proxy.Filter):
+      """This filter causes the proxy to interrupte the download 3 times

[sosa, 2010/12/07 03:44:36]

spelling

[dgarrett, 2010/12/07 04:08:51]

Done.

+      
+         It does this by closing the first three connections to transfer
+         2M total in the outbound connection after they transfer the 
+         2M.
+      """
+      def __init__(self):

[adlr, 2010/12/07 02:26:51]

please add a docstring for each method/function

[dgarrett, 2010/12/07 03:48:18]

Done.

+        # This value is shared across all conections
+        self.close_count = 0
+      
+      def setup(self):
+        # This value is reset for each connection
+        self.data_size = 0
+
+      def OutBound(self, data):
+        # for the first three requests, close the connection after
+        # 2M are sent.
+        if self.close_count < 3:
+          if self.data_size > (2 * 1024 * 1024):
+            self.close_count += 1
+            return None
+        
+        self.data_size += len(data)
+        return data 
+
+
+    self.PrepareBase(target_image_path)
+
+    proxy_port = 8081

[sosa, 2010/12/07 03:44:36]

Comment about the default port choices

[dgarrett, 2010/12/07 04:08:51]

Done.

+    proxy = cros_test_proxy.CrosTestProxy(port_in=proxy_port,
+                                          address_out='127.0.0.1',
+                                          port_out=8080,
+                                          filter=InterruptionFilter())
+    proxy.serve_forever_in_thread()
+
+    # This update is expected to fail...
+    try:
+      self._UpdateImageReportError(target_image_path, proxy_port=proxy_port)
+    finally:
+      proxy.shutdown()
+
+  def testDelayedUpdate(self):
+    """Tests what happens if some data is delayed during update delivery"""
+
+    class DelayedFilter(cros_test_proxy.Filter):
+      """Causes intermittent delays in data transmission.

[sosa, 2010/12/07 03:44:36]

you could probably combine this with the other filter

[dgarrett, 2010/12/07 04:08:51]

I could, but plan to make more changes that will cause them to diverge more.

+      
+         It does this by inserting 3 20 second days when transmitting

[sosa, 2010/12/07 03:44:36]

delays not days :) ... at least i hope

[dgarrett, 2010/12/07 04:08:51]

Done.

+         data after 2M has been sent.
+      """
+      def setup(self):
+        # These values are reset for each connection
+        self.data_size = 0
+        self.delay_count = 0
+
+      def OutBound(self, data):
+        # The first three packets after we reach 2M transferred
+        # are delayed by 20 seconds
+        if self.delay_count < 3:
+          if self.data_size > (2 * 1024 * 1024):
+            self.delay_count += 1
+            time.sleep(20)
+        
+        self.data_size += len(data)
+        return data 
+
+
+    self.PrepareBase(target_image_path)
+
+    proxy_port = 8081
+    proxy = cros_test_proxy.CrosTestProxy(port_in=proxy_port,
+                                          address_out='127.0.0.1',
+                                          port_out=8080,
+                                          filter=DelayedFilter())
+    proxy.serve_forever_in_thread()
+
+    # This update is expected to fail...
+    try:
+      self._UpdateImageReportError(target_image_path, proxy_port=proxy_port)
+    finally:
+      proxy.shutdown()
+
 class RealAUTest(unittest.TestCase, AUTest):
   """Test harness for updating real images."""
 
@@ -266,10 +357,14 @@ class RealAUTest(unittest.TestCase, AUTest):
     """Auto-update to base image to prepare for test."""
     self._UpdateImageReportError(image_path)
 
-  def UpdateImage(self, image_path, stateful_change='old'):
+  def UpdateImage(self, image_path, stateful_change='old', proxy_port=None):
     """Updates a remote image using image_to_live.sh."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
 
+    port_flags = []
+    if proxy_port:
+      port_flags.append('--proxy_port=%s' % proxy_port) 
+
     (code, stdout, stderr) = RunCommandCaptureOutput([
         '%s/image_to_live.sh' % self.crosutils,
         '--image=%s' % image_path,
@@ -277,22 +372,26 @@ class RealAUTest(unittest.TestCase, AUTest):
         stateful_change_flag,
         '--verify',
         '--src_image=%s' % self.source_image
-        ])
+        ] + port_flags)

[sosa, 2010/12/07 03:44:36]

.extend(port_flags)

[dgarrett, 2010/12/07 04:08:51]

Done.

 
     if code != 0:
       raise UpdateException(code, stdout)
 
-  def UpdateUsingPayload(self, update_path, stateful_change='old'):
+  def UpdateUsingPayload(self, update_path, stateful_change='old', proxy_port=None):

[adlr, 2010/12/07 02:26:51]

80 cols

[dgarrett, 2010/12/07 03:48:18]

Done.

     """Updates a remote image using image_to_live.sh."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
 
+    port_flags = []
+    if proxy_port:
+      port_flags.append('--proxy_port=%s' % proxy_port) 
+
     (code, stdout, stderr) = RunCommandCaptureOutput([
         '%s/image_to_live.sh' % self.crosutils,
         '--payload=%s' % update_path,
         '--remote=%s' % remote,
         stateful_change_flag,
         '--verify',
-        ])
+        ] + port_flags)

[sosa, 2010/12/07 03:44:36]

ditto

[dgarrett, 2010/12/07 04:08:51]

Done.

 
     if code != 0:
       raise UpdateException(code, stdout)
@@ -352,12 +451,16 @@ class VirtualAUTest(unittest.TestCase, AUTest):
 
     self.assertTrue(os.path.exists(self.vm_image_path))
 
-  def UpdateImage(self, image_path, stateful_change='old'):
+  def UpdateImage(self, image_path, stateful_change='old', proxy_port=None):
     """Updates VM image with image_path."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
     if self.source_image == base_image_path:
       self.source_image = self.vm_image_path
 
+    port_flags = []

[adlr, 2010/12/07 02:26:51]

there seems to be code duplication here and in the next few functions. can the
common code be combined?

[dgarrett, 2010/12/07 03:48:18]

The changes I've added everywhere are identical, but the code I'm changing isn't
in any of the four cases. It felt like a common function would be as problematic
as leaving them separate.

On 2010/12/07 02:26:51, adlr wrote:

+    if proxy_port:
+      port_flags.append('--proxy_port=%s' % proxy_port) 
+
     (code, stdout, stderr) = RunCommandCaptureOutput([
         '%s/cros_run_vm_update' % self.crosutilsbin,
         '--update_image_path=%s' % image_path,
@@ -368,17 +471,21 @@ class VirtualAUTest(unittest.TestCase, AUTest):
         '--kvm_pid=%s' % _KVM_PID_FILE,
         stateful_change_flag,
         '--src_image=%s' % self.source_image,
-        ])
+        ] + port_flags)
 
     if code != 0:
       raise UpdateException(code, stdout)
 
-  def UpdateUsingPayload(self, update_path, stateful_change='old'):
+  def UpdateUsingPayload(self, update_path, stateful_change='old', proxy_port=None):

[adlr, 2010/12/07 02:26:51]

80 cols

[dgarrett, 2010/12/07 03:48:18]

Done.

     """Updates a remote image using image_to_live.sh."""
     stateful_change_flag = self.GetStatefulChangeFlag(stateful_change)
     if self.source_image == base_image_path:
       self.source_image = self.vm_image_path
 
+    port_flags = []
+    if proxy_port:
+      port_flags.append('--proxy_port=%s' % proxy_port) 
+
     (code, stdout, stderr) = RunCommandCaptureOutput([
         '%s/cros_run_vm_update' % self.crosutilsbin,
         '--payload=%s' % update_path,
@@ -389,7 +496,7 @@ class VirtualAUTest(unittest.TestCase, AUTest):
         '--kvm_pid=%s' % _KVM_PID_FILE,
         stateful_change_flag,
         '--src_image=%s' % self.source_image,
-        ])
+        ] + port_flags)
 
     if code != 0:
       raise UpdateException(code, stdout)