concurrencytest.py 4.9 KB

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