1# Copyright 2020 Google LLC 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14"""Contains convenient helpers for writing tests.""" 15 16import contextlib 17import os 18import sys 19import shutil 20import tempfile 21from unittest import mock 22 23import config_utils 24import docker 25import workspace_utils 26 27INFRA_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 28# pylint: disable=wrong-import-position,import-error 29sys.path.append(INFRA_DIR) 30 31import helper 32 33 34@mock.patch('config_utils._is_dry_run', return_value=True) 35@mock.patch('config_utils.GenericCiEnvironment.project_src_path', 36 return_value=None) 37@mock.patch('os.path.basename', return_value=None) 38def _create_config(config_cls, _, __, ___, **kwargs): 39 """Creates a config object from |config_cls| and then sets every attribute 40 that is a key in |kwargs| to the corresponding value. Asserts that each key in 41 |kwargs| is an attribute of config.""" 42 with mock.patch('config_utils.BaseConfig.validate', return_value=True): 43 config = config_cls() 44 for key, value in kwargs.items(): 45 assert hasattr(config, key), 'Config doesn\'t have attribute: ' + key 46 setattr(config, key, value) 47 48 return config 49 50 51def create_build_config(**kwargs): 52 """Wrapper around _create_config for build configs.""" 53 return _create_config(config_utils.BuildFuzzersConfig, **kwargs) 54 55 56def create_run_config(**kwargs): 57 """Wrapper around _create_config for run configs.""" 58 return _create_config(config_utils.RunFuzzersConfig, **kwargs) 59 60 61def create_workspace(workspace_path='/workspace'): 62 """Returns a workspace located at |workspace_path| ('/workspace' by 63 default).""" 64 config = create_run_config(workspace=workspace_path) 65 return workspace_utils.Workspace(config) 66 67 68def patch_environ(testcase_obj, env=None, empty=False, runner=False): 69 """Patch environment. |testcase_obj| is the unittest.TestCase that contains 70 tests. |env|, if specified, is a dictionary of environment variables to start 71 from. If |empty| is True then the new patched environment will be empty. If 72 |runner| is True then the necessary environment variables will be set to run 73 the scripts from base-runner.""" 74 if env is None: 75 env = {} 76 77 patcher = mock.patch.dict(os.environ, env) 78 testcase_obj.addCleanup(patcher.stop) 79 patcher.start() 80 if empty: 81 for key in os.environ.copy(): 82 del os.environ[key] 83 84 if runner: 85 # Add the scripts for base-runner to the path since the wont be in 86 # /usr/local/bin on host machines during testing. 87 base_runner_dir = os.path.join(INFRA_DIR, 'base-images', 'base-runner') 88 os.environ['PATH'] = (os.environ.get('PATH', '') + os.pathsep + 89 base_runner_dir) 90 if 'GOPATH' not in os.environ: 91 # A GOPATH must be set or else the coverage script fails, even for getting 92 # the coverage of non-Go programs. 93 os.environ['GOPATH'] = '/root/go' 94 95 96@contextlib.contextmanager 97def temp_dir_copy(directory): 98 """Context manager that yields a temporary copy of |directory|.""" 99 with tempfile.TemporaryDirectory() as temp_dir: 100 temp_copy_path = os.path.join(temp_dir, os.path.basename(directory)) 101 shutil.copytree(directory, temp_copy_path) 102 yield temp_copy_path 103 104 105@contextlib.contextmanager 106def docker_temp_dir(): 107 """Returns a temporary a directory that is useful for use with docker. On 108 cleanup this contextmanager uses docker to delete the directory's contents so 109 that if anything is owned by root it can be deleted (which 110 tempfile.TemporaryDirectory() cannot do) by non-root users.""" 111 with tempfile.TemporaryDirectory() as temp_dir: 112 yield temp_dir 113 helper.docker_run([ 114 '-v', f'{temp_dir}:/temp_dir', '-t', docker.BASE_BUILDER_TAG, 115 '/bin/bash', '-c', 'rm -rf /temp_dir/*' 116 ]) 117