test-fixet: Automated VMWare test of AutoUpdate

src/platform/memento_softwareupdate/autoupdated_vm_test.py

Index: src/platform/memento_softwareupdate/autoupdated_vm_test.py
diff --git a/src/platform/memento_softwareupdate/autoupdated_vm_test.py b/src/platform/memento_softwareupdate/autoupdated_vm_test.py
new file mode 100755
index 0000000000000000000000000000000000000000..ce6ebf3681b555ba1c9c6956bb8e5564a5d37d44
--- /dev/null
+++ b/src/platform/memento_softwareupdate/autoupdated_vm_test.py
@@ -0,0 +1,463 @@
+#!/usr/bin/python
+
+# This is a fully automated VM test of autoupdate. Here's what it does:
+# - Downloads the latest dev channel image
+# - Creates a VMware image based on that image
+# - Creates an update for the image under test
+# - Creates a copy of that dev channel image that's one number higher than
+#   the image under test (called the rollback image)
+# - Creates an update for the rollback image
+# - Launches a local HTTP server to pretend to be the AU server
+# - Fires up the VM and waits for it to contact the AU server
+# - AU server gives the VM the image under test update
+# - Waits for the image to be installed, then reboots the VM
+# - AU server gives the VM the rollback image
+# - Waits for the image to be installed, then reboots the VM
+# - Waits for the image to ping the AU server with the rollback image
+# - Done!
+
+# Run this program by passing it a path to the build directory you want
+# to test (this directory should contain rootfs.image).
+
+from xml.dom import minidom
+from BaseHTTPServer import BaseHTTPRequestHandler
+from BaseHTTPServer import HTTPServer
+
+import os
+import re
+import signal
+import socket
+import string
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+
+tmp_dir = '/tmp/au_vm_test'
+scripts_dir = '../../scripts/'
+
+original_version = '0.0.0.0'
+test_version = '0.0.0.0'
+
+# This class stores the state that the server is in and serves as glue
+# between the AU server the control of the VMware process
+class TestState(object):
+  # States we can be in
+  INIT, \
+      INITIAL_WAIT, \
+      TEST_DOWNLOAD, \
+      TEST_WAIT, \
+      ROLLBACK_DOWNLOAD, \
+      ROLLBACK_WAIT = \
+      xrange(6)
+
+  def __init__(self, orig_vers, test_vers, rollback_vers, vm):
+    self.reboot_timeout_seconds = 60 * 4
+    self.orig_vers = orig_vers
+    self.test_vers = test_vers
+    self.rollback_vers = rollback_vers
+    self.SetState(TestState.INIT)
+    self.vm = vm
+  
+  def Die(self, message):
+    print message
+    self.vm.Destroy()
+    # TODO(adlr): exit the entire process, not just this tread
+    sys.exit(1)
+  
+  def SetState(self, state):
+    self.state = state
+  
+  # Should be called to start the VM initially
+  def Start(self):
+    if self.state != TestState.INIT:
+      self.Die('Start called while in bad state')
+    self.SetState(TestState.INITIAL_WAIT)
+    self.vm.Launch()
+    # Kick off timer to wait for the AU ping
+    self.timer = threading.Timer(self.reboot_timeout_seconds,
+                                 self.StartupTimeout)
+    self.timer.start()
+
+  def StartupTimeout(self):
+    self.Die('VM Failed to start and ping' + str(id(self)))
+    # TODO(adlr): exit the entire process, not just this tread
+    sys.exit(1)
+  
+  def FinishInstallTimeout(self):
+    self.vm.Shutdown()
+    self.vm.Launch()
+    self.timer = threading.Timer(60 * 5,  # Sometimes VMWare is very slow
+                                 self.StartupTimeout)
+    self.timer.start()
+
+  # Called by AU server when an update request comes in. Should return
+  # the version that the server should return to the AU client, or
+  # None if no update.
+  def HandleUpdateRequest(self, from_version):
+    print 'HandleUpdateRequest(%s) id:%s state:%s' % \
+      (from_version, id(self), self.state)
+    ret = None
+    # Only cancel timer if we're waiting for the machine to startup
+    if self.timer != None and self.state == TestState.INITIAL_WAIT and \
+       from_version == self.orig_vers:
+      print 'Successfully booted initial'
+      self.timer.cancel()
+      self.timer = None
+    elif self.timer != None and self.state == TestState.TEST_WAIT and \
+         from_version == self.test_vers:
+      print 'Successfully booted test'
+      self.timer.cancel()
+      self.timer = None
+    elif self.timer != None and self.state == TestState.ROLLBACK_WAIT and \
+         from_version == self.rollback_vers:
+      print 'Successfully booted rollback'
+      self.timer.cancel()
+      self.timer = None
+      print 'All done!'
+      # TODO(adlr): exit the entire process, not just this tread
+      sys.exit(0)
+
+    # Pick the version to return
+    if from_version == self.orig_vers:
+      ret = self.test_vers
+    elif from_version == self.test_vers:
+      ret = self.rollback_vers
+
+    # Checks to make sure we move through states correctly
+    if from_version == self.orig_vers:
+      if self.state != TestState.INITIAL_WAIT and \
+         self.state != TestState.TEST_DOWNLOAD and \
+         self.state != TestState.INITIAL_WAIT:
+        self.Die('Error: Request from %s while in state %s' %
+                 (from_version, self.state))
+    elif from_version == self.test_vers:
+      if self.state != TestState.TEST_WAIT and \
+         self.state != TestState.ROLLBACK_DOWNLOAD and \
+         self.state != TestState.ROLLBACK_WAIT:
+        self.Die('Error: Request from %s while in state %s' %
+                 (from_version, self.state))
+    else:
+      print 'odd version to be pinged from: %s' % from_version
+      print 'state is %s' % self.state
+
+    # Update state if needed
+    if self.state == TestState.INITIAL_WAIT:
+      self.SetState(TestState.TEST_DOWNLOAD)
+    elif self.state == TestState.TEST_WAIT:
+      self.SetState(TestState.ROLLBACK_DOWNLOAD)
+
+    if ret is not None:
+      return ret
+    print 'Ignoring update request while in state %s' % self.state
+    return ''
+
+  # Called by AU server when the AU client has finished downloading an image
+  def ImageDownloadComplete(self):
+    print 'ImageDownloadComplete()'
+    valid_state = False
+    if self.state == TestState.TEST_DOWNLOAD:
+      valid_state = True
+      self.SetState(TestState.TEST_WAIT)
+    if self.state == TestState.ROLLBACK_DOWNLOAD:
+      valid_state = True
+      self.SetState(TestState.ROLLBACK_WAIT)
+    if not valid_state:
+      print 'Image download done called at invalid state'
+      # TODO(adlr): exit the entire process, not just this tread
+      sys.exit(1)
+    # Put a timer to reboot the VM
+    if self.timer is not None:
+      self.timer.cancel()
+      self.timer = None
+    self.timer = threading.Timer(self.reboot_timeout_seconds,
+                                 self.FinishInstallTimeout)
+    self.timer.start()
+    return
+
+# This subclass of HTTPServer contains info about the versions of
+# software that the AU server should know about. The AUServerHandler
+# object(s) will access this data.
+class AUHTTPServer(HTTPServer):
+  def __init__(self, ip_port, klass):
+    HTTPServer.__init__(self, ip_port, klass)
+    self.update_info = {}
+    self.files = {}
+  
+  # For a given version of the software, the URL, size, and hash of the update
+  # that gives the user that version of the software.
+  def AddUpdateInfo(self, version, url, size, the_hash):
+    self.update_info[version] = (url, the_hash, size)
+    return
+  
+  # For a given path part of a url, return to the client the file at file_path
+  def AddServedFile(self, url_path, file_path):
+    self.files[url_path] = file_path
+  
+  def SetTestState(self, test_state):
+    self.test_state = test_state
+
+# This class handles HTTP requests. POST requests are when the client
+# is pinging to see if there's an update. GET requests are to download
+# an update.
+class AUServerHandler(BaseHTTPRequestHandler):
+  def do_GET(self):
+    self.send_response(200)
+    self.end_headers()
+    print 'GET: %s' % self.path
+
+    if self.server.files[self.path] != None:
+      print 'GET returning path %s' % self.server.files[self.path]
+      f = open(self.server.files[self.path])
+      while True:
+        data = f.read(1024 * 1024 * 8)
+        if not data:
+          break
+        self.wfile.write(data)
+        self.wfile.flush()
+      f.close()
+      self.server.test_state.ImageDownloadComplete()
+    else:
+      print 'GET returning no path'
+      self.wfile.write(self.path + '\n')
+    return
+
+  def do_POST(self):
+    # Parse the form data posted
+    post_length = int(self.headers.getheader('content-length'))
+    post_data = self.rfile.read(post_length)
+    
+    update_dom = minidom.parseString(post_data)
+    root = update_dom.firstChild
+    query = root.getElementsByTagName('o:app')[0]
+    client_version = query.getAttribute('version')
+    print 'Got update request from %s' % client_version
+
+    # Send response
+    self.send_response(200)
+    self.end_headers()
+
+    new_version = self.server.test_state.HandleUpdateRequest(client_version)
+    print 'Appropriate new version is: %s' % new_version
+
+    if self.server.update_info[new_version] == None:
+      print 'Not sure how to serve reply for %s' % new_version
+      return
+
+    payload = """<?xml version="1.0" encoding="UTF-8"?>
+      <gupdate xmlns="http://www.google.com/update2/response" protocol="2.0">
+        <app appid="{87efface-864d-49a5-9bb3-4b050a7c227a}" status="ok">
+          <ping status="ok"/>
+          <updatecheck
+            codebase="%s"
+            hash="%s"
+            needsadmin="false"
+            size="%s"
+            status="ok"/>
+        </app>
+      </gupdate>
+    """ % self.server.update_info[new_version]
+
+    self.wfile.write(payload)
+    return
+
+# A wrapper for the vmplayer process. Can Launch/Shutdown a vm.
+class vmplayer(object):
+  def __init__(self, filename):
+    self.filename = filename
+  
+  # Launch may (read: probably will) return before the OS has booted
+  def Launch(self):
+    self.process = subprocess.Popen(['/usr/bin/vmplayer', self.filename])
+    self.running = True
+  
+  def Destroy(self):
+    if self.running:
+      self.Shutdown()
+
+  # Shutdown will not return until the vmplayer process has fully terminated
+  # and any cleanup is done.
+  def Shutdown(self):
+    # Pretend user sent Ctrl-C to the vmplayer process
+    os.kill(self.process.pid, signal.SIGINT)
+
+    # Wait while vmplayer saves the vm state...
+    self.process.wait()
+    
+    # Delete the saved vm state
+    # TODO(adlr): remove the state file from the disk
+    self.process = None
+    subprocess.check_call(['/bin/sed', '-i', '/checkpoint.vmState/d',
+                           self.filename])
+    self.running = False
+
+def MakePath(path):
+  subprocess.check_call(['/bin/mkdir', '-p', path])
+
+def DownloadLatestImage(out_path):
+  url = 'http://codf30.jail.google.com/internal/archive/' + \
+        'x86-image-official/LATEST-dev-channel/image.zip' + \
+        '.NOT_SAFE_FOR_USB_INSTALL'
+  url = 'http://www.corp.google.com/~adlr/adlr_test_orig.zip'
+  subprocess.check_call(['/usr/bin/wget', '-O', out_path, url])
+
+def UnzipImage(path, directory):
+  subprocess.check_call(['/usr/bin/unzip', path, '-d', directory])
+
+# Create a stateful partition with a text file that points the AU client
+# to localhost at local_port (which is this very script).
+def CreateDefaultStatefulPartition(local_ip, local_port, out_dir, out_file):
+  # Create sparse file for partition
+  part_size = 512 * 1024 * 1024  # bytes, so 500 MiB
+
+  subprocess.check_call(["""#!/bin/bash
+  set -ex
+  OUT_FILE="%s"
+  STATE_DIR="%s"
+  SIZE="%s"
+  dd if=/dev/zero of="$OUT_FILE" bs=1 count=1 seek=$(($SIZE - 1))
+  mkfs.ext3 -F "$OUT_FILE"
+  mkdir -p "$STATE_DIR"
+  sudo mount -o loop "$OUT_FILE" "$STATE_DIR"
+  sudo mkdir -p "$STATE_DIR/etc"
+  cat <<EOF |sudo dd of="$STATE_DIR/etc/lsb-release"
+CHROMEOS_AUSERVER=http://%s:%s/update
+HTTP_SERVER_OVERRIDE=true
+EOF
+  for i in "$STATE_DIR/etc/lsb-release" "$STATE_DIR/etc" "$STATE_DIR"; do
+    sudo chown root:root "$i"
+    if [ -d "$i" ]; then
+      sudo chmod 0755 "$i"
+    else
+      sudo chmod 0644 "$i"
+    fi
+  done
+  sudo umount -d "$STATE_DIR"
+  """ % (out_file, out_dir + '/state', part_size, local_ip, local_port)],
+    shell=True)
+  return
+  
+
+def CreateVMForImage(image_dir):
+  subprocess.check_call([scripts_dir + 'image_to_vmware.sh', '--from',
+                         image_dir, '--to', image_dir, '--state_image',
+                         image_dir + '/state.image'])
+  return
+
+# Returns (size, hash, path) of the generated image.gz for the given rootfs
+def CreateUpdateForImage(rootfs_image):
+  output = subprocess.Popen([scripts_dir + 'mk_memento_images.sh',
+                             rootfs_image],
+                            stdout=subprocess.PIPE).communicate()[0]
+  matches = re.search('Success. hash is ([^\n]+)', output)
+  the_hash = matches.group(1)
+  path = os.path.dirname(rootfs_image) + '/update.gz'
+  size = os.path.getsize(path)
+  return (size, the_hash, path)
+
+# Modify rootfs 'root_img' to have a new version new_version
+def ModifyImageForRollback(root_img, new_version):
+  subprocess.check_call(["""#!/bin/bash
+  set -ex
+  DIR=$(mktemp -d)
+  sudo mount -o loop "%s" $DIR
+  # update versions in lsb-release
+  sudo sed -i \\
+      -e 's/\\(^GOOGLE_RELEASE=\\|CHROMEOS_RELEASE_VERSION=\\).*/\\1%s/' \\
+      -e 's/^\\(CHROMEOS_RELEASE_DESCRIPTION=.*\\)/\\1-ROLLBACK/' \\
+      "$DIR"/etc/lsb-release
+  sudo umount -d $DIR
+  """ % (root_img, new_version)], shell=True)
+
+# Returns the release version of a rootfs (e.g. 0.6.39.201003241739-a1)
+def GetVersionForRootfs(rootfs_image):
+  mount_dir = tempfile.mkdtemp()
+  subprocess.check_call(['sudo', 'mount', '-o', 'loop,ro',
+                         rootfs_image, mount_dir])
+  version = subprocess.Popen(['awk', '-F', '=',
+                              '/GOOGLE_RELEASE=/{print $2}',
+                              mount_dir + '/etc/lsb-release'],
+                             stdout=subprocess.PIPE).communicate()[0].rstrip()
+  subprocess.check_call(['sudo', 'umount', '-d', mount_dir])
+  subprocess.check_call(['sudo', 'rm', '-rf', mount_dir])
+  return version
+
+# For a given version, increment the last number by 1. E.g.:
+# IncrementVersionNumber('0.23.144.842') = '0.23.144.843'
+def IncrementVersionNumber(version):
+  parts = version.split('.')
+  parts[-1] = str(int(parts[-1]) + 1)
+  return string.join(parts, '.')
+
+def UnpackRootfs(directory):
+  subprocess.check_call(["""#!/bin/bash -x
+    cd "%s"
+    ./unpack_partitions.sh chromiumos_image.bin
+    mv -f part_3 rootfs.image
+    """ % directory], shell=True)
+
+def main():
+  if len(sys.argv) != 2:
+    print 'usage: %s path/to/new/image/dir' % sys.argv[0]
+    sys.exit(1)
+  orig_dir = tmp_dir + '/orig'
+  new_dir = sys.argv[1]
+  rollback_dir = tmp_dir + '/rollback'
+  
+  state_image = orig_dir + '/state.image'
+  port = 8080
+
+  MakePath(tmp_dir)
+
+  # Download latest dev channel release
+  orig_zip = tmp_dir + '/orig.zip'
+  DownloadLatestImage(orig_zip)
+  UnzipImage(orig_zip, orig_dir)
+  UnpackRootfs(orig_dir)
+  orig_version = GetVersionForRootfs(orig_dir + '/rootfs.image')
+  print 'Have original image at version: %s' % orig_version
+
+  # Create new AU image
+  print 'Creating update.gz for test image'
+  UnpackRootfs(new_dir)
+  new_update_details = CreateUpdateForImage(new_dir + '/rootfs.image')
+  new_version = GetVersionForRootfs(new_dir + '/rootfs.image')
+  print 'Have test image at version: %s' % new_version
+
+  # Create rollback image
+  rollback_version = IncrementVersionNumber(new_version)
+  print 'Creating rollback image'
+  UnzipImage(orig_zip, rollback_dir)
+  UnpackRootfs(rollback_dir)
+  ModifyImageForRollback(rollback_dir + '/rootfs.image', rollback_version)
+  print 'Creating update.gz for rollback image'
+  rollback_update_details = CreateUpdateForImage(rollback_dir + '/rootfs.image')
+  print 'Have rollback image at version: %s' % rollback_version
+
+  CreateDefaultStatefulPartition(socket.gethostname(), port, orig_dir,
+                                 state_image)
+  CreateVMForImage(orig_dir)
+  
+  player = vmplayer(orig_dir + '/chromeos.vmx')
+  
+  test_state = TestState(orig_version, new_version, rollback_version, player)
+  
+  server = AUHTTPServer((socket.gethostname(), port), AUServerHandler)
+  
+  base_url = 'http://%s:%s' % (socket.gethostname(), port)
+  
+  server.SetTestState(test_state)
+  server.AddUpdateInfo(new_version, base_url + '/' + new_version,
+                       new_update_details[0], new_update_details[1])
+  server.AddUpdateInfo(rollback_version, base_url + '/' + rollback_version,
+                       rollback_update_details[0], rollback_update_details[1])
+  server.AddServedFile('/' + new_version, new_update_details[2])
+  server.AddServedFile('/' + rollback_version, rollback_update_details[2])
+  
+  test_state.Start()
+  print 'Starting server, use <Ctrl-C> to stop'
+  server.serve_forever()
+
+if __name__ == '__main__':
+  main()