def test_does_ensure_table_created_first(self):
db = DatabaseConnection(":memory:")
cursor = db.execute('PRAGMA table_info(records)')
schema = [col[:3] for col in cursor.fetchall()]
expected_schema = [
(0, 'id', 'TEXT'),
(1, 'request_id', 'TEXT'),
(2, 'source', 'TEXT'),
(3, 'event_type', 'TEXT'),
(4, 'timestamp', 'INTEGER'),
(5, 'payload', 'TEXT'),
]
self.assertEqual(expected_schema, schema)
def test_can_connect_to_argument_file(self, mock_connect):
expected_location = os.path.expanduser(
os.path.join('~', 'foo', 'bar', 'baz.db'))
DatabaseConnection(expected_location)
mock_connect.assert_called_with(expected_location,
check_same_thread=False,
isolation_level=None)
def attach_history_handler(session, parsed_args, **kwargs):
if _should_enable_cli_history(session, parsed_args):
LOG.debug('Enabling CLI history')
history_filename = os.environ.get(HISTORY_FILENAME_ENV_VAR,
DEFAULT_HISTORY_FILENAME)
if not os.path.isdir(os.path.dirname(history_filename)):
os.makedirs(os.path.dirname(history_filename))
connection = DatabaseConnection(history_filename)
writer = DatabaseRecordWriter(connection)
record_builder = RecordBuilder()
db_handler = DatabaseHistoryHandler(writer, record_builder)
HISTORY_RECORDER.add_handler(db_handler)
HISTORY_RECORDER.enable()
def setUp(self):
self.connection = DatabaseConnection(':memory:')
self.connection.row_factory = sqlite3.Row
def setUp(self):
self.db = DatabaseConnection(':memory:')
self.writer = DatabaseRecordWriter(connection=self.db)
self.record_builder = RecordBuilder()
self.handler = DatabaseHistoryHandler(
writer=self.writer, record_builder=self.record_builder)
class TestDatabaseHistoryHandler(unittest.TestCase):
UUID_PATTERN = re.compile(
'^[0-9a-f]{8}-([0-9a-f]{4}-){3}[0-9a-f]{12}$',
re.I
)
def setUp(self):
self.db = DatabaseConnection(':memory:')
self.writer = DatabaseRecordWriter(connection=self.db)
self.record_builder = RecordBuilder()
self.handler = DatabaseHistoryHandler(
writer=self.writer, record_builder=self.record_builder)
def _get_last_record(self):
record = self.db.execute('SELECT * FROM records').fetchone()
return record
def _assert_expected_event_type(self, source, record):
self.assertEqual(source, record[3])
def _assert_expected_payload(self, source, record):
loaded_payload = json.loads(record[-1])
self.assertEqual(source, loaded_payload)
def _assert_expected_source(self, source, record):
self.assertEqual(source, record[2])
def _assert_has_request_id(self, record):
identifier = record[1]
self.assertTrue(self.UUID_PATTERN.match(identifier))
def _assert_record_has_command_id(self, record):
identifier = record[0]
self.assertTrue(self.UUID_PATTERN.match(identifier))
def test_does_emit_write_record(self):
self.handler.emit('event_type', 'payload', 'source')
record = self._get_last_record()
self._assert_record_has_command_id(record)
self._assert_expected_event_type('event_type', record)
self._assert_expected_payload('payload', record)
self._assert_expected_source('source', record)
def test_can_emit_write_record_with_structure(self):
payload = {'foo': 'bar'}
self.handler.emit('event_type', payload, 'source')
record = self._get_last_record()
self._assert_record_has_command_id(record)
self._assert_expected_event_type('event_type', record)
self._assert_expected_payload(payload, record)
self._assert_expected_source('source', record)
def test_can_emit_cli_version_record(self):
# CLI_VERSION records have a list of strings payload
payload = 'foobarbaz'
self.handler.emit('CLI_VERSION', payload, 'CLI')
record = self._get_last_record()
self._assert_record_has_command_id(record)
self._assert_expected_event_type('CLI_VERSION', record)
self._assert_expected_payload(payload, record)
self._assert_expected_source('CLI', record)
def test_can_emit_cli_arguments_record(self):
# CLI_ARGUMENTS records have a list of strings payload
payload = ['foo', 'bar', 'baz']
self.handler.emit('CLI_ARGUMENTS', payload, 'CLI')
record = self._get_last_record()
self._assert_record_has_command_id(record)
self._assert_expected_event_type('CLI_ARGUMENTS', record)
self._assert_expected_payload(payload, record)
self._assert_expected_source('CLI', record)
def test_can_emit_api_call_record(self):
# API_CALL records have a dictionary based payload
payload = {
'service': 's3',
'operation': 'ListBuckets',
'params': {}
}
self.handler.emit('API_CALL', payload, 'BOTOCORE')
record = self._get_last_record()
self._assert_record_has_command_id(record)
self._assert_has_request_id(record)
self._assert_expected_event_type('API_CALL', record)
self._assert_expected_payload(payload, record)
self._assert_expected_source('BOTOCORE', record)
def test_can_emit_api_call_record_with_binary_param(self):
# API_CALL records have a dictionary based payload
payload = {
'service': 'lambda',
'operation': 'CreateFunction',
'params': {
"FunctionName": "Name",
"Handler": "mod.fn",
"Role": "foobar",
"Runtime": "python3",
"Code": {
"ZipFile": b'zipfile binary content \xfe\xed'
}
}
}
self.handler.emit('API_CALL', payload, 'BOTOCORE')
record = self._get_last_record()
parsed_payload = payload.copy()
parsed_payload['params']['Code']['ZipFile'] = \
'<Byte sequence>'
self._assert_record_has_command_id(record)
self._assert_has_request_id(record)
self._assert_expected_event_type('API_CALL', record)
self._assert_expected_payload(parsed_payload, record)
self._assert_expected_source('BOTOCORE', record)
def test_can_emit_http_request_record(self):
# HTTP_REQUEST records have have their entire body field as a binary
# blob, howver it will all be utf-8 valid since the binary fields
# from the api call will have been b64 encoded.
payload = {
'url': ('https://lambda.us-west-2.amazonaws.com/2015-03-31/'
'functions'),
'method': 'POST',
'headers': CaseInsensitiveDict({
'foo': 'bar'
}),
'body': b'body with no invalid utf-8 bytes in it',
'streaming': False
}
self.handler.emit('HTTP_REQUEST', payload, 'BOTOCORE')
record = self._get_last_record()
parsed_payload = payload.copy()
parsed_payload['headers'] = dict(parsed_payload['headers'])
parsed_payload['body'] = 'body with no invalid utf-8 bytes in it'
self._assert_record_has_command_id(record)
self._assert_expected_event_type('HTTP_REQUEST', record)
self._assert_expected_payload(parsed_payload, record)
self._assert_expected_source('BOTOCORE', record)
def test_can_emit_http_response_record(self):
# HTTP_RESPONSE also contains a binary response in its body, but it
# will not contain any non-unicode characters
payload = {
'status_code': 200,
'headers': CaseInsensitiveDict({
'foo': 'bar'
}),
'body': b'body with no invalid utf-8 bytes in it',
'streaming': False
}
self.handler.emit('HTTP_RESPONSE', payload, 'BOTOCORE')
record = self._get_last_record()
parsed_payload = payload.copy()
parsed_payload['headers'] = dict(parsed_payload['headers'])
parsed_payload['body'] = 'body with no invalid utf-8 bytes in it'
self._assert_record_has_command_id(record)
self._assert_expected_event_type('HTTP_RESPONSE', record)
self._assert_expected_payload(parsed_payload, record)
self._assert_expected_source('BOTOCORE', record)
def test_can_emit_parsed_response_record(self):
payload = {
"Count": 1,
"Items": [
{
"strkey": {
"S": "string"
}
}
],
"ScannedCount": 1,
"ConsumedCapacity": None
}
self.handler.emit('PARSED_RESPONSE', payload, 'BOTOCORE')
record = self._get_last_record()
self._assert_record_has_command_id(record)
self._assert_expected_event_type('PARSED_RESPONSE', record)
self._assert_expected_payload(payload, record)
self._assert_expected_source('BOTOCORE', record)
def test_can_emit_parsed_response_record_with_binary(self):
# PARSED_RESPONSE can also contain raw bytes
payload = {
"Count": 1,
"Items": [
{
"bitkey": {
"B": b"binary data \xfe\xed"
}
}
],
"ScannedCount": 1,
"ConsumedCapacity": None
}
self.handler.emit('PARSED_RESPONSE', payload, 'BOTOCORE')
record = self._get_last_record()
parsed_payload = payload.copy()
parsed_payload['Items'][0]['bitkey']['B'] = "<Byte sequence>"
self._assert_record_has_command_id(record)
self._assert_expected_event_type('PARSED_RESPONSE', record)
self._assert_expected_payload(payload, record)
self._assert_expected_source('BOTOCORE', record)
def test_does_not_mutate_dict(self):
payload = {
"bitkey": b"binary data \xfe\xed"
}
copy_payload = payload.copy()
self.handler.emit('test', payload, 'BOTOCORE')
self.assertEqual(payload, copy_payload)
def test_does_not_mutate_list(self):
payload = ['non binary data', b"binary data \xfe\xed"]
copy_payload = list(payload)
self.handler.emit('test', payload, 'BOTOCORE')
self.assertEqual(payload, copy_payload)
def test_can_close(self, mock_connect):
connection = mock.Mock()
mock_connect.return_value = connection
conn = DatabaseConnection(':memory:')
conn.close()
self.assertTrue(connection.close.called)
def test_does_try_to_enable_wal(self, mock_connect):
conn = DatabaseConnection(':memory:')
conn._connection.execute.assert_any_call('PRAGMA journal_mode=WAL')
def setUp(self):
self.db = DatabaseConnection(':memory:')
self.writer = DatabaseRecordWriter(self.db)
def _connect_to_history_db(self):
if self._db_reader is None:
connection = DatabaseConnection(self._get_history_db_filename())
self._db_reader = DatabaseRecordReader(connection)