Index: scripts/slave/unittests/expect_tests/handle_train.py |
diff --git a/scripts/slave/unittests/expect_tests/handle_train.py b/scripts/slave/unittests/expect_tests/handle_train.py |
new file mode 100644 |
index 0000000000000000000000000000000000000000..b918130abfaad2d1ca44c7f6ad05e43627b5b2d0 |
--- /dev/null |
+++ b/scripts/slave/unittests/expect_tests/handle_train.py |
@@ -0,0 +1,109 @@ |
+# Copyright 2014 The Chromium Authors. All rights reserved. |
+# Use of this source code is governed by a BSD-style license that can be |
+# found in the LICENSE file. |
+ |
+import time |
+import os |
+import collections |
+ |
+from .types import Handler |
+from .serialize import WriteNewData, DiffData, NonExistant, GetCurrentData |
+ |
+ |
+DirSeen = collections.namedtuple('DirSeen', 'dir') |
+WriteAction = collections.namedtuple('WriteAction', 'test why') |
+NoAction = collections.namedtuple('NoAction', 'test') |
+ |
+ |
+class TrainHandler(Handler): |
+ """Write test expectations to disk.""" |
+ @classmethod |
+ def add_options(cls, parser): |
+ parser.add_argument( |
+ '--force', action='store_true', help=( |
+ 'Immediately write expectations to disk instead of determining if ' |
+ 'they contain a diff from the current expectations.' |
+ )) |
+ |
+ @classmethod |
+ def gen_stage_loop(cls, _opts, tests, put_next_stage, put_result_stage): |
+ dirs_seen = set() |
+ for test in tests: |
+ if test.expect_dir not in dirs_seen: |
+ try: |
+ os.makedirs(test.expect_dir) |
+ except OSError: |
+ pass |
+ put_result_stage(DirSeen(test.expect_dir)) |
+ dirs_seen.add(test.expect_dir) |
+ put_next_stage(test) |
+ |
+ @classmethod |
+ def run_stage_loop(cls, opts, tests_results, put_next_stage): |
+ for test, result in tests_results: |
+ if opts.force: |
+ WriteNewData(test, result.data) |
+ put_next_stage(WriteAction(test, 'forced')) |
+ return |
+ |
+ current, same_schema = GetCurrentData(test) |
+ diff = DiffData(current, result.data) |
+ if diff is not None or not same_schema: |
+ WriteNewData(test, result.data) |
+ if current is NonExistant: |
+ why = 'missing' |
+ elif diff: |
+ why = 'diff' |
+ else: |
+ why = 'schema changed' |
+ put_next_stage(WriteAction(test, why)) |
+ else: |
+ put_next_stage(NoAction(test)) |
+ |
+ class ResultStageHandler(Handler.ResultStageHandler): |
+ def __init__(self, opts): |
+ super(TrainHandler.ResultStageHandler, self).__init__(opts) |
+ self.dirs_seen = set() |
+ self.files_expected = collections.defaultdict(set) |
+ self.start = time.time() |
+ self.num_tests = 0 |
+ |
+ def _record_expected(self, test): |
+ head, tail = os.path.split(test.expect_path()) |
+ self.files_expected[head].add(tail) |
+ |
+ # Handlers |
+ def DirSeen(self, dirseen): |
+ self.dirs_seen.add(dirseen.dir) |
+ |
+ def NoAction(self, result): |
+ self._record_expected(result.test) |
+ if self.opts.verbose: |
+ print '%s did not change' % result.test.name |
+ |
+ def WriteAction(self, result): |
+ self._record_expected(result.test) |
+ if not self.opts.quiet: |
+ test = result.test |
+ name = test.expect_path() if self.opts.verbose else test.name |
+ print 'Wrote %s: %s' % (name, result.why) |
+ |
+ def finalize(self, aborted): |
+ if not aborted and not self.opts.test_glob: |
+ for d in self.dirs_seen: |
+ expected = self.files_expected[d] |
+ for f in os.listdir(d): |
+ if f == 'OWNERS': |
+ continue |
+ if f not in expected: |
+ path = os.path.join(d, f) |
+ os.unlink(path) |
+ if self.opts.verbose: |
+ print 'Removed unexpected file', path |
+ if not self.opts.quiet: |
+ num_tests = sum(len(x) for x in self.files_expected.itervalues()) |
+ print 'Trained %d tests in %0.3fs' % ( |
+ num_tests, time.time() - self.start) |
+ |
+ |
+ |