def test_main_aggregate(self, mock_reader): mock_reader.return_value = [SAMPLE_RECORDS[0], SAMPLE_RECORDS[0]] with patch('sys.stdout', new_callable=StringIO) as mock_stdout: main(['mygroup', 'aggregate']) output = mock_stdout.getvalue().splitlines() actual_header = output[0].split('\t') expected_header = [ 'bytes', 'dstaddr', 'dstport', 'end', 'packets', 'protocol', 'srcaddr', 'srcport', 'start', ] self.assertEqual(actual_header, expected_header) actual_line = output[1].split('\t') expected_line = [ '1680', '192.0.2.1', '49152', '2015-08-12 13:47:44', '20', '6', '198.51.100.1', '443', '2015-08-12 13:47:43', ] self.assertEqual(actual_line, expected_line)
def test_main_assume_role_s3(self, mock_boto3, mock_reader): mock_boto3.client.return_value.assume_role.return_value = { 'Credentials': { 'AccessKeyId': 'myaccesskeyid', 'SecretAccessKey': 'mysecretaccesskey', 'SessionToken': 'mysessiontoken', } } mock_client = MagicMock() mock_boto3.session.Session.return_value.client.return_value = ( mock_client ) mock_reader.return_value = [] args = [ '--role-arn', 'myarn', '--external-id', 'uuid4', '--location-type', 's3', 'mybucket' ] main(args) session = mock_boto3.session.Session session.assert_called_once_with( aws_access_key_id='myaccesskeyid', aws_secret_access_key='mysecretaccesskey', aws_session_token='mysessiontoken', ) session.return_value.client.assert_called_once_with('s3') mock_reader.assert_called_once_with( 'mybucket', boto_client=mock_client )
def test_main(self, mock_reader): main(['mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', region_name='us-east-1' ) main(['-s', '2015-05-05 14:20:00', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', region_name='us-east-1', start_time=datetime(2015, 5, 5, 14, 20) ) main(['--end-time', '2015-05-05 14:20:00', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', region_name='us-east-1', end_time=datetime(2015, 5, 5, 14, 20) ) main([ '--time-format', '%Y-%m-%d', '--start-time', '2015-05-05', 'mygroup' ]) mock_reader.assert_called_with( log_group_name='mygroup', region_name='us-east-1', start_time=datetime(2015, 5, 5) ) main(['--region', 'us-west-1', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', region_name='us-west-1' )
def test_main_assume_role(self, mock_boto3, mock_reader): mock_boto3.client.return_value.assume_role.return_value = { 'Credentials': { 'AccessKeyId': 'myaccesskeyid', 'SecretAccessKey': 'mysecretaccesskey', 'SessionToken': 'mysessiontoken', } } mock_client = MagicMock() mock_boto3.session.Session.return_value.client.return_value = ( mock_client ) mock_reader.return_value = [] main(['--role-arn', 'myarn', '--external-id', 'uuid4', 'mygroup']) session = mock_boto3.session.Session session.assert_called_once_with( aws_access_key_id='myaccesskeyid', aws_secret_access_key='mysecretaccesskey', aws_session_token='mysessiontoken', ) session.return_value.client.assert_called_once_with('logs') mock_reader.assert_called_once_with( log_group_name='mygroup', boto_client=mock_client )
def test_main_print(self, mock_out, mock_reader): mock_out.stdout = io.BytesIO() mock_reader.return_value = SAMPLE_RECORDS main(['mygroup']) for call, record in zip_longest(mock_out.mock_calls, SAMPLE_INPUT): __, args, kwargs = call line = args[0] self.assertEqual(line, record)
def test_main_findip(self, mock_out, mock_reader): mock_out.stdout = io.BytesIO() mock_reader.return_value = SAMPLE_RECORDS main(['mygroup', 'findip', '198.51.100.2']) expected_result = [SAMPLE_INPUT[2]] for call, record in zip_longest(mock_out.mock_calls, expected_result): __, args, kwargs = call line = args[0] self.assertEqual(line, record)
def test_main_missing_arn(self, mock_out, mock_reader): mock_out.stdout = io.BytesIO() mock_reader.return_value = SAMPLE_RECORDS main(['--external-id', 'uuid4', 'mygroup']) expected_result = [ 'must give a --role-arn if an --external-id is given', ] for call, result in zip_longest(mock_out.mock_calls, expected_result): __, args, kwargs = call line = args[0] self.assertEqual(line, result)
def test_main_bad_action(self, mock_out, mock_reader): mock_out.stdout = io.BytesIO() mock_reader.return_value = SAMPLE_RECORDS main(['mygroup', '__']) expected_result = [ 'unknown action: __', 'known actions: {}'.format(', '.join(actions)), ] for call, result in zip_longest(mock_out.mock_calls, expected_result): __, args, kwargs = call line = args[0] self.assertEqual(line, result)
def test_s3_destination(self, mock_out, mock_reader): mock_out.stdout = io.BytesIO() mock_reader.return_value = SAMPLE_RECORDS main( [ 'mybucket/myprefix', '--location-type', 's3', '--include-accounts', '999999999998, 999999999999', '--include-regions', 'us-east-1,us-east-2', ] ) mock_reader.assert_called_once_with( location='mybucket/myprefix', include_accounts=['999999999998', '999999999999'], include_regions=['us-east-1', 'us-east-2'], ) for call, record in zip_longest(mock_out.mock_calls, SAMPLE_INPUT): __, args, kwargs = call line = args[0] self.assertEqual(line, record)
def test_main_ipset(self, mock_out, mock_reader): mock_out.stdout = io.BytesIO() mock_reader.return_value = SAMPLE_RECORDS main(['mygroup', 'ipset']) expected_set = set() for record in SAMPLE_INPUT: data = record.split() expected_set.add(data[3]) expected_set.add(data[4]) # don't include SKIPDATA/NODATA in results expected_set.remove('-') # make sure the number of lines are the same as the size of the set self.assertEqual(len(mock_out.mock_calls), len(expected_set)) actual_set = set() for __, args, kwargs in mock_out.mock_calls: line = args[0] actual_set.add(line) self.assertEqual(actual_set, expected_set)
def test_main_assume_role(self, mock_boto3, mock_reader): mock_boto3.client.return_value.assume_role.return_value = { 'Credentials': { 'AccessKeyId': 'myaccesskeyid', 'SecretAccessKey': 'mysecretaccesskey', 'SessionToken': 'mysessiontoken', } } mock_client = MagicMock() mock_boto3.session.Session.return_value.client.return_value = ( mock_client ) mock_reader.return_value = [] main(['--role-arn', 'myarn', '--external-id', 'uuid4', 'mygroup']) mock_boto3.session.Session.assert_called_once_with( aws_access_key_id='myaccesskeyid', aws_secret_access_key='mysecretaccesskey', aws_session_token='mysessiontoken', ) mock_reader.assert_called_once_with( log_group_name='mygroup', boto_client=mock_client )
def test_main_print_count(self, mock_out, mock_reader): mock_out.stdout = io.BytesIO() mock_reader.return_value = SAMPLE_RECORDS with self.assertRaises(ValueError): main(['mygroup', 'print', 'two']) with self.assertRaises(RuntimeError): main(['mygroup', 'print', '2', '3']) main(['mygroup', 'print', '2']) for call, record in zip_longest(mock_out.mock_calls, SAMPLE_INPUT[:2]): __, args, kwargs = call line = args[0] self.assertEqual(line, record)
def test_main(self, mock_reader): main(['mygroup']) mock_reader.assert_called_with(log_group_name='mygroup') main(['-s', '2015-05-05 14:20:00', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', start_time=datetime(2015, 5, 5, 14, 20), ) main(['--end-time', '2015-05-05 14:20:00', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', end_time=datetime(2015, 5, 5, 14, 20), ) main([ '--time-format', '%Y-%m-%d', '--start-time', '2015-05-05', 'mygroup' ]) mock_reader.assert_called_with( log_group_name='mygroup', start_time=datetime(2015, 5, 5), ) main(['--region', 'us-west-1', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', region_name='us-west-1', ) main(['--profile', 'my-profile', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', profile_name='my-profile' ) main(['--filter-pattern', 'REJECT', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', filter_pattern='REJECT' )
def test_main(self, mock_reader): main(['mygroup']) mock_reader.assert_called_with(log_group_name='mygroup', fields=None) main(['-s', '2015-05-05 14:20:00', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', fields=None, start_time=datetime(2015, 5, 5, 14, 20), ) main(['--end-time', '2015-05-05 14:20:00', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', fields=None, end_time=datetime(2015, 5, 5, 14, 20), ) main([ '--time-format', '%Y-%m-%d', '--start-time', '2015-05-05', 'mygroup', '--thread-count', '2', ]) mock_reader.assert_called_with( log_group_name='mygroup', fields=None, start_time=datetime(2015, 5, 5), thread_count=2, ) main(['--region', 'us-west-1', 'mygroup']) mock_reader.assert_called_with( log_group_name='mygroup', fields=None, region_name='us-west-1', ) main(['--profile', 'my-profile', 'mygroup']) mock_reader.assert_called_with(log_group_name='mygroup', fields=None, profile_name='my-profile') main([ '--filter-pattern', 'REJECT', '--fields', '${account-id} ${action}', 'mygroup', ]) mock_reader.assert_called_with( log_group_name='mygroup', fields=('account-id', 'action'), filter_pattern='REJECT', )