def test_predict_num_roles(): """ Test predict function of HostFootprint class with varying number of distinct roles present """ with tempfile.TemporaryDirectory() as tmpdir: testdata = os.path.join(tmpdir, 'test_data') shutil.copytree('./tests/test_data', testdata) for file in ['combined_three_roles.csv', 'combined_two_roles.csv']: input_file = os.path.join(testdata, file) operation = 'train' sys.argv = hf_args(tmpdir, operation, input_file) instance = HostFootprint() instance.main() operation = 'predict' sys.argv = hf_args(tmpdir, operation, input_file) instance = HostFootprint() instance.main() predictions = json.loads(instance.predict()) assert isinstance(predictions, dict) # Check if number of predictions is correct if file == 'combined_three_roles.csv': assert len(predictions) == 6 else: assert len(predictions) == 4
def test_train(): """Test training function of HostFootprint class""" with tempfile.TemporaryDirectory() as tmpdir: testdata = os.path.join(tmpdir, 'test_data') shutil.copytree('./tests/test_data', testdata) input_file = os.path.join(testdata, 'combined.csv') operation = 'train' sys.argv = hf_args(tmpdir, operation, input_file) instance = HostFootprint() instance.main()
def test_train_bad_data_too_few_columns(): """ This test tries to train a model on a mal-formed csv with too few fields """ with tempfile.TemporaryDirectory() as tmpdir: testdata = os.path.join(tmpdir, 'test_data') shutil.copytree('./tests/test_data', testdata) input_file = os.path.join(testdata, 'bad_data_too_few_columns.csv') operation = 'train' sys.argv = hf_args(tmpdir, operation, input_file) instance = HostFootprint() with pytest.raises(Exception): instance.main()
def run_algorithm_stage(self, in_path): raw_args = self.add_opt_args(self.stage_args['algorithm']) raw_args.extend(['-O', self.operation, '-v', self.log_level, in_path]) instance = HostFootprint(raw_args=raw_args) return instance.main()