1*760c253cSXin Li# Copyright 2019 The ChromiumOS Authors 2*760c253cSXin Li# Use of this source code is governed by a BSD-style license that can be 3*760c253cSXin Li# found in the LICENSE file. 4*760c253cSXin Li 5*760c253cSXin Li"""Helper functions for unit testing.""" 6*760c253cSXin Li 7*760c253cSXin Liimport contextlib 8*760c253cSXin Liimport json 9*760c253cSXin Liimport os 10*760c253cSXin Liimport tempfile 11*760c253cSXin Li 12*760c253cSXin Li 13*760c253cSXin Liclass ArgsOutputTest: 14*760c253cSXin Li """Testing class to simulate a argument parser object.""" 15*760c253cSXin Li 16*760c253cSXin Li def __init__(self, svn_option="google3"): 17*760c253cSXin Li self.chromeos_path = "/abs/path/to/chroot" 18*760c253cSXin Li self.last_tested = "/abs/path/to/last_tested_file.json" 19*760c253cSXin Li self.llvm_version = svn_option 20*760c253cSXin Li self.extra_change_lists = None 21*760c253cSXin Li self.options = ["latest-toolchain"] 22*760c253cSXin Li self.builders = ["some-builder"] 23*760c253cSXin Li 24*760c253cSXin Li 25*760c253cSXin Li# FIXME: Migrate modules with similar helper to use this module. 26*760c253cSXin Lidef CallCountsToMockFunctions(mock_function): 27*760c253cSXin Li """A decorator that passes a call count to the function it decorates. 28*760c253cSXin Li 29*760c253cSXin Li Examples: 30*760c253cSXin Li @CallCountsToMockFunctions 31*760c253cSXin Li def foo(call_count): 32*760c253cSXin Li return call_count 33*760c253cSXin Li ... 34*760c253cSXin Li ... 35*760c253cSXin Li [foo(), foo(), foo()] 36*760c253cSXin Li [0, 1, 2] 37*760c253cSXin Li """ 38*760c253cSXin Li 39*760c253cSXin Li counter = [0] 40*760c253cSXin Li 41*760c253cSXin Li def Result(*args, **kwargs): 42*760c253cSXin Li # For some values of `counter`, the mock function would simulate 43*760c253cSXin Li # raising an exception, so let the test case catch the exception via 44*760c253cSXin Li # `unittest.TestCase.assertRaises()` and to also handle recursive 45*760c253cSXin Li # functions. 46*760c253cSXin Li prev_counter = counter[0] 47*760c253cSXin Li counter[0] += 1 48*760c253cSXin Li 49*760c253cSXin Li ret_value = mock_function(prev_counter, *args, **kwargs) 50*760c253cSXin Li 51*760c253cSXin Li return ret_value 52*760c253cSXin Li 53*760c253cSXin Li return Result 54*760c253cSXin Li 55*760c253cSXin Li 56*760c253cSXin Lidef WritePrettyJsonFile(file_name, json_object): 57*760c253cSXin Li """Writes the contents of the file to the json object. 58*760c253cSXin Li 59*760c253cSXin Li Args: 60*760c253cSXin Li file_name: The file that has contents to be used for the json object. 61*760c253cSXin Li json_object: The json object to write to. 62*760c253cSXin Li """ 63*760c253cSXin Li 64*760c253cSXin Li json.dump(file_name, json_object, indent=4, separators=(",", ": ")) 65*760c253cSXin Li 66*760c253cSXin Li 67*760c253cSXin Lidef CreateTemporaryJsonFile(): 68*760c253cSXin Li """Makes a temporary .json file.""" 69*760c253cSXin Li 70*760c253cSXin Li return CreateTemporaryFile(suffix=".json") 71*760c253cSXin Li 72*760c253cSXin Li 73*760c253cSXin Li@contextlib.contextmanager 74*760c253cSXin Lidef CreateTemporaryFile(suffix=""): 75*760c253cSXin Li """Makes a temporary file.""" 76*760c253cSXin Li 77*760c253cSXin Li fd, temp_file_path = tempfile.mkstemp(suffix=suffix) 78*760c253cSXin Li 79*760c253cSXin Li os.close(fd) 80*760c253cSXin Li 81*760c253cSXin Li try: 82*760c253cSXin Li yield temp_file_path 83*760c253cSXin Li 84*760c253cSXin Li finally: 85*760c253cSXin Li if os.path.isfile(temp_file_path): 86*760c253cSXin Li os.remove(temp_file_path) 87