utils/python/ctftestsuite/stress.py

   1 import argparse
   2 import shutil
   3 import sys
   4 import os
   5
   6
   7 class MetadataTest:
   8     def __init__(self):
   9         self._actions = {
  10             'prepare': self._prepare,
  11             'clean': self._clean,
  12         }
  13
  14     @staticmethod
  15     def _perror(str):
  16         print('Error: {}'.format(str), file=sys.stderr)
  17         sys.exit(1)
  18
  19     def _parse_args(self):
  20         ap = argparse.ArgumentParser()
  21
  22         ap.add_argument('-s', '--size', action='store', metavar='SIZE',
  23                         type=int, required=True, help='size')
  24         ap.add_argument('action', metavar='ACTION', action='store',
  25                         help='action')
  26
  27         # parse args
  28         args = ap.parse_args()
  29
  30         # validate size
  31         if args.size < 1:
  32             MetadataTest._perror('wrong size: {}'.format(args.size))
  33
  34         return args
  35
  36     def _get_what(self):
  37         return self.what.format(size=self.size)
  38
  39     def _prepare(self):
  40         # make sure everything is clean first
  41         self._clean()
  42
  43         print('Preparing test for {}'.format(self._get_what()))
  44
  45         # make test directory
  46         os.mkdir(self._trace_dir_path)
  47
  48         # open and write metadata file
  49         with open(self._metadata_path, 'w') as f:
  50             self.write_metadata(f)
  51
  52     def _clean(self):
  53         print('Cleaning up test for {}'.format(self._get_what()))
  54
  55         try:
  56             shutil.rmtree(self._trace_dir_path, ignore_errors=True)
  57         except Exception as e:
  58             # ignore
  59             pass
  60
  61     def _do_action(self):
  62         if self._action not in self._actions:
  63             MetadataTest._perror('invalid action: "{}"'.format(self._action))
  64
  65         self._actions[self._action]()
  66
  67     @property
  68     def size(self):
  69         return self._size
  70
  71     def main(self):
  72         args = self._parse_args()
  73         self._size = args.size
  74         self._trace_dir_path = 'trace-{}'.format(args.size)
  75         self._metadata_path = os.path.join(self._trace_dir_path, 'metadata')
  76         self._action = args.action
  77
  78         self._do_action()