1*11ae93eeSSimon Glass#!/usr/bin/env python 2*11ae93eeSSimon Glass# SPDX-License-Identifier: GPL-2.0+ 3*11ae93eeSSimon Glass# 4*11ae93eeSSimon Glass# Modified by: Corey Goldberg, 2013 5*11ae93eeSSimon Glass# 6*11ae93eeSSimon Glass# Original code from: 7*11ae93eeSSimon Glass# Bazaar (bzrlib.tests.__init__.py, v2.6, copied Jun 01 2013) 8*11ae93eeSSimon Glass# Copyright (C) 2005-2011 Canonical Ltd 9*11ae93eeSSimon Glass 10*11ae93eeSSimon Glass"""Python testtools extension for running unittest suites concurrently. 11*11ae93eeSSimon Glass 12*11ae93eeSSimon GlassThe `testtools` project provides a ConcurrentTestSuite class, but does 13*11ae93eeSSimon Glassnot provide a `make_tests` implementation needed to use it. 14*11ae93eeSSimon Glass 15*11ae93eeSSimon GlassThis allows you to parallelize a test run across a configurable number 16*11ae93eeSSimon Glassof worker processes. While this can speed up CPU-bound test runs, it is 17*11ae93eeSSimon Glassmainly useful for IO-bound tests that spend most of their time waiting for 18*11ae93eeSSimon Glassdata to arrive from someplace else and can benefit from cocncurrency. 19*11ae93eeSSimon Glass 20*11ae93eeSSimon GlassUnix only. 21*11ae93eeSSimon Glass""" 22*11ae93eeSSimon Glass 23*11ae93eeSSimon Glassimport os 24*11ae93eeSSimon Glassimport sys 25*11ae93eeSSimon Glassimport traceback 26*11ae93eeSSimon Glassimport unittest 27*11ae93eeSSimon Glassfrom itertools import cycle 28*11ae93eeSSimon Glassfrom multiprocessing import cpu_count 29*11ae93eeSSimon Glass 30*11ae93eeSSimon Glassfrom subunit import ProtocolTestCase, TestProtocolClient 31*11ae93eeSSimon Glassfrom subunit.test_results import AutoTimingTestResultDecorator 32*11ae93eeSSimon Glass 33*11ae93eeSSimon Glassfrom testtools import ConcurrentTestSuite, iterate_tests 34*11ae93eeSSimon Glass 35*11ae93eeSSimon Glass 36*11ae93eeSSimon Glass_all__ = [ 37*11ae93eeSSimon Glass 'ConcurrentTestSuite', 38*11ae93eeSSimon Glass 'fork_for_tests', 39*11ae93eeSSimon Glass 'partition_tests', 40*11ae93eeSSimon Glass] 41*11ae93eeSSimon Glass 42*11ae93eeSSimon Glass 43*11ae93eeSSimon GlassCPU_COUNT = cpu_count() 44*11ae93eeSSimon Glass 45*11ae93eeSSimon Glass 46*11ae93eeSSimon Glassdef fork_for_tests(concurrency_num=CPU_COUNT): 47*11ae93eeSSimon Glass """Implementation of `make_tests` used to construct `ConcurrentTestSuite`. 48*11ae93eeSSimon Glass 49*11ae93eeSSimon Glass :param concurrency_num: number of processes to use. 50*11ae93eeSSimon Glass """ 51*11ae93eeSSimon Glass def do_fork(suite): 52*11ae93eeSSimon Glass """Take suite and start up multiple runners by forking (Unix only). 53*11ae93eeSSimon Glass 54*11ae93eeSSimon Glass :param suite: TestSuite object. 55*11ae93eeSSimon Glass 56*11ae93eeSSimon Glass :return: An iterable of TestCase-like objects which can each have 57*11ae93eeSSimon Glass run(result) called on them to feed tests to result. 58*11ae93eeSSimon Glass """ 59*11ae93eeSSimon Glass result = [] 60*11ae93eeSSimon Glass test_blocks = partition_tests(suite, concurrency_num) 61*11ae93eeSSimon Glass # Clear the tests from the original suite so it doesn't keep them alive 62*11ae93eeSSimon Glass suite._tests[:] = [] 63*11ae93eeSSimon Glass for process_tests in test_blocks: 64*11ae93eeSSimon Glass process_suite = unittest.TestSuite(process_tests) 65*11ae93eeSSimon Glass # Also clear each split list so new suite has only reference 66*11ae93eeSSimon Glass process_tests[:] = [] 67*11ae93eeSSimon Glass c2pread, c2pwrite = os.pipe() 68*11ae93eeSSimon Glass pid = os.fork() 69*11ae93eeSSimon Glass if pid == 0: 70*11ae93eeSSimon Glass try: 71*11ae93eeSSimon Glass stream = os.fdopen(c2pwrite, 'wb', 1) 72*11ae93eeSSimon Glass os.close(c2pread) 73*11ae93eeSSimon Glass # Leave stderr and stdout open so we can see test noise 74*11ae93eeSSimon Glass # Close stdin so that the child goes away if it decides to 75*11ae93eeSSimon Glass # read from stdin (otherwise its a roulette to see what 76*11ae93eeSSimon Glass # child actually gets keystrokes for pdb etc). 77*11ae93eeSSimon Glass sys.stdin.close() 78*11ae93eeSSimon Glass subunit_result = AutoTimingTestResultDecorator( 79*11ae93eeSSimon Glass TestProtocolClient(stream) 80*11ae93eeSSimon Glass ) 81*11ae93eeSSimon Glass process_suite.run(subunit_result) 82*11ae93eeSSimon Glass except: 83*11ae93eeSSimon Glass # Try and report traceback on stream, but exit with error 84*11ae93eeSSimon Glass # even if stream couldn't be created or something else 85*11ae93eeSSimon Glass # goes wrong. The traceback is formatted to a string and 86*11ae93eeSSimon Glass # written in one go to avoid interleaving lines from 87*11ae93eeSSimon Glass # multiple failing children. 88*11ae93eeSSimon Glass try: 89*11ae93eeSSimon Glass stream.write(traceback.format_exc()) 90*11ae93eeSSimon Glass finally: 91*11ae93eeSSimon Glass os._exit(1) 92*11ae93eeSSimon Glass os._exit(0) 93*11ae93eeSSimon Glass else: 94*11ae93eeSSimon Glass os.close(c2pwrite) 95*11ae93eeSSimon Glass stream = os.fdopen(c2pread, 'rb', 1) 96*11ae93eeSSimon Glass test = ProtocolTestCase(stream) 97*11ae93eeSSimon Glass result.append(test) 98*11ae93eeSSimon Glass return result 99*11ae93eeSSimon Glass return do_fork 100*11ae93eeSSimon Glass 101*11ae93eeSSimon Glass 102*11ae93eeSSimon Glassdef partition_tests(suite, count): 103*11ae93eeSSimon Glass """Partition suite into count lists of tests.""" 104*11ae93eeSSimon Glass # This just assigns tests in a round-robin fashion. On one hand this 105*11ae93eeSSimon Glass # splits up blocks of related tests that might run faster if they shared 106*11ae93eeSSimon Glass # resources, but on the other it avoids assigning blocks of slow tests to 107*11ae93eeSSimon Glass # just one partition. So the slowest partition shouldn't be much slower 108*11ae93eeSSimon Glass # than the fastest. 109*11ae93eeSSimon Glass partitions = [list() for _ in range(count)] 110*11ae93eeSSimon Glass tests = iterate_tests(suite) 111*11ae93eeSSimon Glass for partition, test in zip(cycle(partitions), tests): 112*11ae93eeSSimon Glass partition.append(test) 113*11ae93eeSSimon Glass return partitions 114*11ae93eeSSimon Glass 115*11ae93eeSSimon Glass 116*11ae93eeSSimon Glassif __name__ == '__main__': 117*11ae93eeSSimon Glass import time 118*11ae93eeSSimon Glass 119*11ae93eeSSimon Glass class SampleTestCase(unittest.TestCase): 120*11ae93eeSSimon Glass """Dummy tests that sleep for demo.""" 121*11ae93eeSSimon Glass 122*11ae93eeSSimon Glass def test_me_1(self): 123*11ae93eeSSimon Glass time.sleep(0.5) 124*11ae93eeSSimon Glass 125*11ae93eeSSimon Glass def test_me_2(self): 126*11ae93eeSSimon Glass time.sleep(0.5) 127*11ae93eeSSimon Glass 128*11ae93eeSSimon Glass def test_me_3(self): 129*11ae93eeSSimon Glass time.sleep(0.5) 130*11ae93eeSSimon Glass 131*11ae93eeSSimon Glass def test_me_4(self): 132*11ae93eeSSimon Glass time.sleep(0.5) 133*11ae93eeSSimon Glass 134*11ae93eeSSimon Glass # Load tests from SampleTestCase defined above 135*11ae93eeSSimon Glass suite = unittest.TestLoader().loadTestsFromTestCase(SampleTestCase) 136*11ae93eeSSimon Glass runner = unittest.TextTestRunner() 137*11ae93eeSSimon Glass 138*11ae93eeSSimon Glass # Run tests sequentially 139*11ae93eeSSimon Glass runner.run(suite) 140*11ae93eeSSimon Glass 141*11ae93eeSSimon Glass # Run same tests across 4 processes 142*11ae93eeSSimon Glass suite = unittest.TestLoader().loadTestsFromTestCase(SampleTestCase) 143*11ae93eeSSimon Glass concurrent_suite = ConcurrentTestSuite(suite, fork_for_tests(4)) 144*11ae93eeSSimon Glass runner.run(concurrent_suite) 145