testing/legion/examples/http_example/http_test.py

   1 # Copyright 2015 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 import argparse
   6 import logging
   7 import os
   8 import socket
   9 import sys
  10 import time
  11
  12 TESTING_DIR = os.path.join(
  13     os.path.dirname(os.path.abspath(__file__)), '..', '..', '..')
  14 sys.path.append(TESTING_DIR)
  15
  16 from legion import legion_test_case
  17
  18
  19 class HttpTest(legion_test_case.TestCase):
  20   """Example HTTP test case."""
  21
  22   @classmethod
  23   def GetArgs(cls):
  24     """Get command line args."""
  25     parser = argparse.ArgumentParser()
  26     parser.add_argument('--http-server')
  27     parser.add_argument('--http-client')
  28     parser.add_argument('--os', default='Ubuntu-14.04')
  29     args, _ = parser.parse_known_args()
  30     return args
  31
  32   @classmethod
  33   def CreateTask(cls, name, task_hash, os_type):
  34     """Create a new task."""
  35     #pylint: disable=unexpected-keyword-arg,no-value-for-parameter
  36     #pylint: disable=arguments-differ
  37     task = super(HttpTest, cls).CreateTask(
  38         name=name,
  39         isolated_hash=task_hash,
  40         dimensions={'os': os_type})
  41     task.Create()
  42     return task
  43
  44   @classmethod
  45   def setUpClass(cls):
  46     """Creates the task machines and waits until they connect."""
  47     args = cls.GetArgs()
  48     cls.http_server = cls.CreateTask(
  49         'http_server', args.http_server, args.os)
  50     cls.http_client = cls.CreateTask(
  51         'http_client', args.http_client, args.os)
  52     cls.http_server.WaitForConnection()
  53     cls.http_client.WaitForConnection()
  54
  55   def CanConnectToServerPort(self, server_port):
  56     """Connect to a port on the http_server.
  57
  58     Returns:
  59       True if the connection succeeded, False otherwise.
  60     """
  61     try:
  62       s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  63       s.connect((self.http_server.ip_address, server_port))
  64       return True
  65     except socket.error:
  66       return False
  67
  68   def FindOpenPortOnServer(self):
  69     """Find an open port on the server and return it.
  70
  71     Returns:
  72       The value of an open port.
  73     """
  74     for server_port in  xrange(2000, 20000):
  75       if not self.CanConnectToServerPort(server_port):
  76         return server_port
  77     self.fail('Unable to find an open port on the server.')
  78
  79   def StartServer(self, server_port):
  80     """Starts the http_server process.
  81
  82     Returns:
  83       The server process.
  84     """
  85     def WaitForServer():
  86       timeout = time.time() + 5
  87       while timeout > time.time():
  88         if self.CanConnectToServerPort(server_port):
  89           return
  90       self.fail('Server process failed to start')
  91
  92     cmd = [
  93         self.http_server.executable,
  94         'http_server.py',
  95         '--port', str(server_port)
  96         ]
  97     proc = self.http_server.Process(cmd)
  98     WaitForServer()
  99     return proc
 100
 101   def StartClient(self, server_port):
 102     """Starts the http_client process.
 103
 104     Returns:
 105       The client process.
 106     """
 107     cmd = [
 108         self.http_client.executable,
 109         'http_client.py',
 110         '--server', self.http_server.ip_address,
 111         '--port', str(server_port)
 112         ]
 113     return self.http_client.Process(cmd)
 114
 115   def testHttpWorks(self):
 116     """Tests that the client process can talk to the server process."""
 117     server_proc = None
 118     client_proc = None
 119     try:
 120       server_port = self.FindOpenPortOnServer()
 121       logging.info('Starting server at %s:%s', self.http_server.ip_address,
 122                    server_port)
 123       server_proc = self.StartServer(server_port)
 124       logging.info('Connecting to server at %s:%s', self.http_server.ip_address,
 125                    server_port)
 126       client_proc = self.StartClient(server_port)
 127       client_proc.Wait()
 128       logging.info('client_proc.stdout: %s', client_proc.ReadStdout())
 129       logging.info('client_proc.stderr: %s', client_proc.ReadStderr())
 130       self.assertEqual(client_proc.GetReturncode(), 0)
 131     finally:
 132       if server_proc:
 133         server_proc.Kill()
 134         server_proc.Delete()
 135       if client_proc:
 136         client_proc.Kill()
 137         client_proc.Delete()
 138
 139
 140 if __name__ == '__main__':
 141   legion_test_case.main()