test_corpus.py 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. #!/usr/bin/env python3
  2. #
  3. # Copyright 2018 The Chromium OS Authors. All rights reserved.
  4. # Use of this source code is governed by a BSD-style license that can be
  5. # found in the LICENSE file.
  6. #
  7. """A tool for running Puffin tests in a corpus of deflate compressed files."""
  8. import argparse
  9. import filecmp
  10. import logging
  11. import os
  12. import subprocess
  13. import sys
  14. import tempfile
  15. _PUFFHUFF = 'puffhuff'
  16. _PUFFDIFF = 'puffdiff'
  17. TESTS = (_PUFFHUFF, _PUFFDIFF)
  18. class Error(Exception):
  19. """Puffin general processing error."""
  20. def ParseArguments(argv):
  21. """Parses and Validates command line arguments.
  22. Args:
  23. argv: command line arguments to parse.
  24. Returns:
  25. The arguments list.
  26. """
  27. parser = argparse.ArgumentParser()
  28. parser.add_argument('corpus', metavar='CORPUS',
  29. help='A corpus directory containing compressed files')
  30. parser.add_argument('-d', '--disabled_tests', default=(), metavar='',
  31. nargs='*',
  32. help=('Space separated list of tests to disable. '
  33. 'Allowed options include: ' + ', '.join(TESTS)),
  34. choices=TESTS)
  35. parser.add_argument('--cache_size', type=int, metavar='SIZE',
  36. help='The size (in bytes) of the cache for puffpatch '
  37. 'operations.')
  38. parser.add_argument('--debug', action='store_true',
  39. help='Turns on verbosity.')
  40. # Parse command-line arguments.
  41. args = parser.parse_args(argv)
  42. if not os.path.isdir(args.corpus):
  43. raise Error('Corpus directory {} is non-existent or inaccesible'
  44. .format(args.corpus))
  45. return args
  46. def main(argv):
  47. """The main function."""
  48. args = ParseArguments(argv[1:])
  49. if args.debug:
  50. logging.getLogger().setLevel(logging.DEBUG)
  51. # Construct list of appropriate files.
  52. files = list(filter(os.path.isfile, [os.path.join(args.corpus, f)
  53. for f in os.listdir(args.corpus)]))
  54. # For each file in corpus run puffhuff.
  55. if _PUFFHUFF not in args.disabled_tests:
  56. for src in files:
  57. with tempfile.NamedTemporaryFile() as tgt_file:
  58. operation = 'puffhuff'
  59. logging.debug('Running %s on %s', operation, src)
  60. cmd = ['puffin',
  61. '--operation={}'.format(operation),
  62. '--src_file={}'.format(src),
  63. '--dst_file={}'.format(tgt_file.name)]
  64. if subprocess.call(cmd) != 0:
  65. raise Error('Puffin failed to do {} command: {}'
  66. .format(operation, cmd))
  67. if not filecmp.cmp(src, tgt_file.name):
  68. raise Error('The generated file {} is not equivalent to the original '
  69. 'file {} after {} operation.'
  70. .format(tgt_file.name, src, operation))
  71. if _PUFFDIFF not in args.disabled_tests:
  72. # Run puffdiff and puffpatch for each pairs of files in the corpus.
  73. for src in files:
  74. for tgt in files:
  75. with tempfile.NamedTemporaryFile() as patch, \
  76. tempfile.NamedTemporaryFile() as new_tgt:
  77. operation = 'puffdiff'
  78. logging.debug('Running %s on %s (%d) and %s (%d)',
  79. operation,
  80. os.path.basename(src), os.stat(src).st_size,
  81. os.path.basename(tgt), os.stat(tgt).st_size)
  82. cmd = ['puffin',
  83. '--operation={}'.format(operation),
  84. '--src_file={}'.format(src),
  85. '--dst_file={}'.format(tgt),
  86. '--patch_file={}'.format(patch.name)]
  87. # Running the puffdiff operation
  88. if subprocess.call(cmd) != 0:
  89. raise Error('Puffin failed to do {} command: {}'
  90. .format(operation, cmd))
  91. logging.debug('Patch size is: %d', os.stat(patch.name).st_size)
  92. operation = 'puffpatch'
  93. logging.debug('Running %s on src file %s and patch %s',
  94. operation, os.path.basename(src), patch.name)
  95. cmd = ['puffin',
  96. '--operation={}'.format(operation),
  97. '--src_file={}'.format(src),
  98. '--dst_file={}'.format(new_tgt.name),
  99. '--patch_file={}'.format(patch.name)]
  100. if args.cache_size:
  101. cmd += ['--cache_size={}'.format(args.cache_size)]
  102. # Running the puffpatch operation
  103. if subprocess.call(cmd) != 0:
  104. raise Error('Puffin failed to do {} command: {}'
  105. .format(operation, cmd))
  106. if not filecmp.cmp(tgt, new_tgt.name):
  107. raise Error('The generated file {} is not equivalent to the '
  108. 'original file {} after puffpatch operation.'
  109. .format(new_tgt.name, tgt))
  110. return 0
  111. if __name__ == '__main__':
  112. sys.exit(main(sys.argv))