testing/legion/examples/hello_world/controller_test.py

   1 #!/usr/bin/env python
   2 # Copyright 2015 The Chromium Authors. All rights reserved.
   3 # Use of this source code is governed by a BSD-style license that can be
   4 # found in the LICENSE file.
   5
   6 """A simple host test module.
   7
   8 This module runs on the host machine and is responsible for creating 2
   9 task machines, waiting for them, and running RPC calls on them.
  10 """
  11
  12 import argparse
  13 import logging
  14 import os
  15 import sys
  16 import time
  17
  18 # Map the testing directory so we can import legion.legion_test.
  19 TESTING_DIR = os.path.join(
  20     os.path.dirname(os.path.abspath(__file__)),
  21     '..', '..', '..', '..', 'testing')
  22 sys.path.append(TESTING_DIR)
  23
  24 from legion import legion_test_case
  25
  26
  27 class ExampleTestController(legion_test_case.TestCase):
  28   """A simple example controller for a test."""
  29
  30   @classmethod
  31   def CreateTestTask(cls):
  32     """Create a new task."""
  33     parser = argparse.ArgumentParser()
  34     parser.add_argument('--task-hash')
  35     parser.add_argument('--os', default='Ubuntu-14.04')
  36     args, _ = parser.parse_known_args()
  37
  38     task = cls.CreateTask(
  39         isolated_hash=args.task_hash,
  40         dimensions={'os': args.os},
  41         idle_timeout_secs=90,
  42         connection_timeout_secs=90,
  43         verbosity=logging.DEBUG)
  44     task.Create()
  45     return task
  46
  47   @classmethod
  48   def setUpClass(cls):
  49     """Creates the task machines and waits until they connect."""
  50     cls.task1 = cls.CreateTestTask()
  51     cls.task2 = cls.CreateTestTask()
  52     cls.task1.WaitForConnection()
  53     cls.task2.WaitForConnection()
  54
  55   def testCallEcho(self):
  56     """Tests rpc.Echo on a task."""
  57     logging.info('Calling Echo on %s', self.task2.name)
  58     self.assertEqual(self.task2.rpc.Echo('foo'), 'echo foo')
  59
  60   def testLaunchTaskBinary(self):
  61     """Call task_test.py 'name' on the tasks."""
  62     self.VerifyTaskBinaryLaunched(self.task1)
  63     self.VerifyTaskBinaryLaunched(self.task2)
  64
  65   def VerifyTaskBinaryLaunched(self, task):
  66     logging.info(
  67         'Calling Process to run "task_test.py %s"', task.name)
  68     proc = task.Process(['python', 'task_test.py', task.name])
  69     proc.Wait()
  70     self.assertEqual(proc.GetReturncode(), 0)
  71     self.assertIn(task.name, proc.ReadStdout())
  72     self.assertEquals(proc.ReadStderr(), '')
  73     proc.Delete()
  74
  75
  76 if __name__ == '__main__':
  77   legion_test_case.main()