class TaskManager_IntegrationTestCase(unittest.TestCase): TEST_CONFIG = TEST_TASK_CLIENT_CONFIG LABOURER = Labourer( id='some_function', arn='arn:aws:lambda:us-west-2:000000000000:function:some_function') @classmethod def setUpClass(cls): """ Clean the classic autotest table. """ cls.TEST_CONFIG['init_clients'] = ['DynamoDb'] def setUp(self): """ We keep copies of main parameters here, because they may differ from test to test and cleanup needs them. This is responsibility of the test author to update these values if required from test. """ self.config = self.TEST_CONFIG.copy() self.HASH_KEY = ('task_id', 'S') self.RANGE_KEY = ('labourer_id', 'S') self.table_name = self.config['dynamo_db_config']['table_name'] self.completed_tasks_table = self.config['sosw_closed_tasks_table'] self.retry_tasks_table = self.config['sosw_retry_tasks_table'] self.clean_task_tables() self.dynamo_client = DynamoDbClient( config=self.config['dynamo_db_config']) self.manager = TaskManager(custom_config=self.config) self.manager.ecology_client = MagicMock() self.labourer = deepcopy(self.LABOURER) def tearDown(self): self.clean_task_tables() def clean_task_tables(self): clean_dynamo_table(self.table_name, (self.HASH_KEY[0], )) clean_dynamo_table(self.completed_tasks_table, ('task_id', )) clean_dynamo_table(self.retry_tasks_table, ('labourer_id', 'task_id')) def setup_tasks(self, status='available', mutiple_labourers=False, count_tasks=3): """ Some fake adding some scheduled tasks for some workers. """ _ = self.manager.get_db_field_name _cfg = self.manager.config.get table = _cfg('dynamo_db_config')['table_name'] if status not in ['closed', 'failed'] \ else _cfg('sosw_closed_tasks_table') MAP = { 'available': { self.RANGE_KEY[0]: lambda x: str(worker_id), _('greenfield'): lambda x: round(1000 + random.randrange(0, 100000, 1000)), _('attempts'): lambda x: 0, }, 'invoked': { self.RANGE_KEY[0]: lambda x: str(worker_id), _('greenfield'): lambda x: round(time.time()) + _cfg( 'greenfield_invocation_delta'), _('attempts'): lambda x: 1, }, 'expired': { self.RANGE_KEY[0]: lambda x: str(worker_id), _('greenfield'): lambda x: round(time.time()) + _cfg( 'greenfield_invocation_delta') - random.randint( 1000, 10000), _('attempts'): lambda x: 1, }, 'running': { self.RANGE_KEY[0]: lambda x: str(worker_id), _('greenfield'): lambda x: round(time.time()) + _cfg( 'greenfield_invocation_delta') - random.randint(1, 900), _('attempts'): lambda x: 1, }, 'closed': { _('greenfield'): lambda x: round(time.time()) + _cfg( 'greenfield_invocation_delta') - random.randint( 1000, 10000), _('labourer_id_task_status'): lambda x: f"{self.LABOURER.id}_1", _('completed_at'): lambda x: x[_('greenfield')] - _cfg( 'greenfield_invocation_delta') + random.randint(10, 300), _('closed_at'): lambda x: x[_('completed_at')] + random.randint(1, 60), _('attempts'): lambda x: 3, }, 'failed': { _('greenfield'): lambda x: round(time.time()) + _cfg( 'greenfield_invocation_delta') - random.randint( 1000, 10000), _('labourer_id_task_status'): lambda x: f"{self.LABOURER.id}_0", _('closed_at'): lambda x: x[_('greenfield')] + 900 + random.randint(1, 60), _('attempts'): lambda x: 3, }, } # raise ValueError(f"Unsupported `status`: {status}. Should be one of: 'available', 'invoked'.") workers = [self.LABOURER.id] if not mutiple_labourers else range( 42, 45) for worker_id in workers: for i in range(count_tasks): row = { self.HASH_KEY[0]: f"task_id_{worker_id}_{i}_{str(uuid.uuid4())[:8]}", # Task ID } for field, getter in MAP[status].items(): row[field] = getter(row) print(f"Putting {row} to {table}") self.dynamo_client.put(row, table_name=table) time.sleep( 0.1 ) # Sleep a little to fit the Write Capacity (10 WCU) of autotest table. def test_get_next_for_labourer(self): self.setup_tasks() # time.sleep(5) result = self.manager.get_next_for_labourer(self.LABOURER, only_ids=True) # print(result) self.assertEqual(len(result), 1, "Returned more than one task") self.assertIn(f'task_id_{self.LABOURER.id}_', result[0]) def test_get_next_for_labourer__multiple(self): self.setup_tasks() result = self.manager.get_next_for_labourer(self.LABOURER, cnt=5000, only_ids=True) # print(result) self.assertEqual(len(result), 3, "Should be just 3 tasks for this worker in setup") self.assertTrue( all(f'task_id_{self.LABOURER.id}_' in task for task in result), "Returned some tasks of other Workers") def test_get_next_for_labourer__not_take_invoked(self): self.setup_tasks() self.setup_tasks(status='invoked') result = self.manager.get_next_for_labourer(self.LABOURER, cnt=50, only_ids=True) # print(result) self.assertEqual( len(result), 3, "Should be just 3 tasks for this worker in setup. The other 3 are invoked." ) self.assertTrue( all(f'task_id_{self.LABOURER.id}_' in task for task in result), "Returned some tasks of other Workers") def test_get_next_for_labourer__full_tasks(self): self.setup_tasks() result = self.manager.get_next_for_labourer(self.LABOURER, cnt=2) # print(result) self.assertEqual(len(result), 2, "Should be just 2 tasks as requested") for task in result: self.assertIn( f'task_id_{self.LABOURER.id}_', task['task_id']), "Returned some tasks of other Workers" self.assertEqual( self.LABOURER.id, task['labourer_id']), "Returned some tasks of other Workers" def register_labourers(self): self.manager.get_labourers = MagicMock(return_value=[self.LABOURER]) return self.manager.register_labourers() def test_mark_task_invoked(self): greenfield = round(time.time() - random.randint(100, 1000)) delta = self.manager.config['greenfield_invocation_delta'] self.register_labourers() row = { self.HASH_KEY[0]: f"task_id_{self.LABOURER.id}_256", # Task ID self.RANGE_KEY[0]: self.LABOURER.id, # Worker ID 'greenfield': greenfield } self.dynamo_client.put(row) # print(f"Saved initial version with greenfield some date not long ago: {row}") # Do the actual tested job self.manager.mark_task_invoked(self.LABOURER, row) time.sleep(1) result = self.dynamo_client.get_by_query( {self.HASH_KEY[0]: f"task_id_{self.LABOURER.id}_256"}, strict=False) # print(f"The new updated value of task is: {result}") # Rounded -2 we check that the greenfield was updated self.assertAlmostEqual(round(int(time.time()) + delta, -2), round(result[0]['greenfield'], -2)) def test_get_invoked_tasks_for_labourer(self): self.register_labourers() self.setup_tasks(status='running') self.setup_tasks(status='expired') self.setup_tasks(status='invoked') self.assertEqual( len(self.manager.get_invoked_tasks_for_labourer(self.LABOURER)), 3) def test_get_running_tasks_for_labourer(self): self.register_labourers() self.setup_tasks(status='available') self.setup_tasks(status='running') self.setup_tasks(status='expired') self.assertEqual( len(self.manager.get_running_tasks_for_labourer(self.LABOURER)), 3) def test_get_expired_tasks_for_labourer(self): self.register_labourers() self.setup_tasks(status='running') self.setup_tasks(status='expired') self.assertEqual( len(self.manager.get_expired_tasks_for_labourer(self.LABOURER)), 3) # @unittest.skip("Function currently depricated") # def test_close_task(self): # _ = self.manager.get_db_field_name # # Create task with id=123 # task = {_('task_id'): '123', _('labourer_id'): 'lambda1', _('greenfield'): 8888, _('attempts'): 2, # _('completed_at'): 123123} # self.dynamo_client.put(task) # # # Call # self.manager.close_task(task_id='123', labourer_id='lambda1') # # # Get from db, check # tasks = self.dynamo_client.get_by_query({_('task_id'): '123'}) # self.assertEqual(len(tasks), 1) # task_result = tasks[0] # # expected_result = task.copy() # # for k in ['task_id', 'labourer_id', 'greenfield', 'attempts']: # assert expected_result[k] == task_result[k] # # self.assertTrue(_('closed_at') in task_result, msg=f"{_('closed_at')} not in task_result {task_result}") # self.assertTrue(time.time() - 360 < task_result[_('closed_at')] < time.time()) def test_archive_task(self): _ = self.manager.get_db_field_name # Create task with id=123 task = { _('task_id'): '123', _('labourer_id'): 'lambda1', _('greenfield'): 8888, _('attempts'): 2 } self.dynamo_client.put(task) # Call self.manager.archive_task('123') # Check the task isn't in the tasks db, but is in the completed_tasks table tasks = self.dynamo_client.get_by_query({_('task_id'): '123'}) self.assertEqual(len(tasks), 0) completed_tasks = self.dynamo_client.get_by_query( {_('task_id'): '123'}, table_name=self.completed_tasks_table) self.assertEqual(len(completed_tasks), 1) completed_task = completed_tasks[0] for k in task.keys(): self.assertEqual(task[k], completed_task[k]) for k in completed_task.keys(): if k != _('closed_at'): self.assertEqual(task[k], completed_task[k]) self.assertTrue( time.time() - 360 < completed_task[_('closed_at')] < time.time()) def test_move_task_to_retry_table(self): _ = self.manager.get_db_field_name labourer_id = 'lambda1' task = { _('task_id'): '123', _('labourer_id'): labourer_id, _('greenfield'): 8888, _('attempts'): 2 } delay = 300 self.dynamo_client.put(task) self.manager.move_task_to_retry_table(task, delay) result_tasks = self.dynamo_client.get_by_query({_('task_id'): '123'}) self.assertEqual(len(result_tasks), 0) result_retry_tasks = self.dynamo_client.get_by_query( {_('labourer_id'): labourer_id}, table_name=self.retry_tasks_table) self.assertEqual(len(result_retry_tasks), 1) result = first_or_none(result_retry_tasks) for k in task: self.assertEqual(task[k], result[k]) for k in result: if k != _('desired_launch_time'): self.assertEqual(result[k], task[k]) self.assertTrue( time.time() + delay - 60 < result[_('desired_launch_time')] < time.time() + delay + 60) def test_get_tasks_to_retry_for_labourer(self): _ = self.manager.get_db_field_name tasks = RETRY_TASKS.copy() # Add tasks to retry table for task in tasks: self.dynamo_client.put(task, self.config['sosw_retry_tasks_table']) # Call with patch('time.time') as t: t.return_value = 9500 labourer = self.manager.register_labourers()[0] result_tasks = self.manager.get_tasks_to_retry_for_labourer(labourer, limit=20) self.assertEqual(len(result_tasks), 2) # Check it only gets tasks with timestamp <= now self.assertIn(tasks[0], result_tasks) self.assertIn(tasks[1], result_tasks) def test_retry_tasks(self): _ = self.manager.get_db_field_name with patch('time.time') as t: t.return_value = 9500 labourer = self.manager.register_labourers()[0] self.manager.get_oldest_greenfield_for_labourer = Mock( return_value=8888) # Add tasks to tasks_table regular_tasks = [ { _('labourer_id'): labourer.id, _('task_id'): '11', _('arn'): 'some_arn', _('payload'): {}, _('greenfield'): 8888 }, { _('labourer_id'): labourer.id, _('task_id'): '22', _('arn'): 'some_arn', _('payload'): {}, _('greenfield'): 9999 }, ] for task in regular_tasks: self.dynamo_client.put(task) # Add tasks to retry_table retry_tasks = RETRY_TASKS.copy() for task in retry_tasks: self.dynamo_client.put( task, table_name=self.config['sosw_retry_tasks_table']) retry_table_items = self.dynamo_client.get_by_scan( table_name=self.retry_tasks_table) self.assertEqual(len(retry_table_items), len(retry_tasks)) # Use get_tasks_to_retry_for_labourer to get tasks tasks = self.manager.get_tasks_to_retry_for_labourer(labourer) # Call self.manager.retry_tasks(labourer, tasks) # Check removed 2 out of 3 tasks from retry queue. One is desired to be launched later. retry_table_items = self.dynamo_client.get_by_scan( table_name=self.retry_tasks_table) self.assertEqual(len(retry_table_items), 1) # Check tasks moved to `tasks_table` with lowest greenfields tasks_table_items = self.dynamo_client.get_by_scan() for x in tasks_table_items: print(x) self.assertEqual(len(tasks_table_items), 4) for reg_task in regular_tasks: self.assertIn(reg_task, tasks_table_items) for retry_task in retry_tasks: try: matching = next(x for x in tasks_table_items if x[_('task_id')] == retry_task[_('task_id')]) except StopIteration: print( f"Task not retried {retry_task}. Probably not yet desired." ) continue for k in retry_task.keys(): if k not in [_('greenfield'), _('desired_launch_time')]: self.assertEqual(retry_task[k], matching[k]) for k in matching.keys(): if k != _('greenfield'): self.assertEqual(retry_task[k], matching[k]) print( f"New greenfield of a retried task: {matching[_('greenfield')]}" ) self.assertTrue(matching[_('greenfield')] < min( [x[_('greenfield')] for x in regular_tasks])) @patch.object(boto3, '__version__', return_value='1.9.53') def test_retry_tasks__old_boto(self, n): self.test_retry_tasks() def test_get_oldest_greenfield_for_labourer__get_newest_greenfield_for_labourer( self): with patch('time.time') as t: t.return_value = 9500 labourer = self.manager.register_labourers()[0] min_gf = 20000 max_gf = 10000 for i in range(5): # Ran this with range(1000), it passes :) gf = random.randint(10000, 20000) if gf < min_gf: min_gf = gf if gf > max_gf: max_gf = gf row = { 'labourer_id': f"{labourer.id}", 'task_id': f"task-{i}", 'greenfield': gf } self.dynamo_client.put(row) time.sleep( 0.1 ) # Sleep a little to fit the Write Capacity (10 WCU) of autotest table. result = self.manager.get_oldest_greenfield_for_labourer(labourer) self.assertEqual(min_gf, result) newest = self.manager.get_newest_greenfield_for_labourer(labourer) self.assertEqual(max_gf, newest) def test_get_length_of_queue_for_labourer(self): labourer = Labourer(id='some_lambda', arn='some_arn') num_of_tasks = 3 # Ran this with 464 tasks and it worked for i in range(num_of_tasks): row = { 'labourer_id': f"some_lambda", 'task_id': f"task-{i}", 'greenfield': i } self.dynamo_client.put(row) time.sleep( 0.1 ) # Sleep a little to fit the Write Capacity (10 WCU) of autotest table. queue_len = self.manager.get_length_of_queue_for_labourer(labourer) self.assertEqual(queue_len, num_of_tasks) def test_get_average_labourer_duration__calculates_average__only_failing_tasks( self): self.manager.ecology_client.get_max_labourer_duration.return_value = 900 some_labourer = self.register_labourers()[0] self.setup_tasks(status='failed', count_tasks=15) self.assertEqual( 900, self.manager.get_average_labourer_duration(some_labourer)) def test_get_average_labourer_duration__calculates_average(self): self.manager.ecology_client.get_max_labourer_duration.return_value = 900 some_labourer = self.register_labourers()[0] self.setup_tasks(status='closed', count_tasks=15) self.setup_tasks(status='failed', count_tasks=15) self.assertLessEqual( self.manager.get_average_labourer_duration(some_labourer), 900) self.assertGreaterEqual( self.manager.get_average_labourer_duration(some_labourer), 10)
class DynamodbClientIntegrationTestCase(unittest.TestCase): TEST_CONFIG = { 'row_mapper': { 'lambda_name': 'S', 'invocation_id': 'S', 'en_time': 'N', 'hash_col': 'S', 'range_col': 'N', 'other_col': 'S', 'new_col': 'S', 'some_col': 'S', 'some_counter': 'N', 'some_bool': 'BOOL', 'some_map': 'M', }, 'required_fields': ['lambda_name'], 'table_name': 'autotest_dynamo_db', 'hash_key': 'hash_col' } @classmethod def setUpClass(cls): clean_dynamo_table() def setUp(self): self.HASH_COL = 'hash_col' self.HASH_KEY = (self.HASH_COL, 'S') self.RANGE_COL = 'range_col' self.RANGE_COL_TYPE = 'N' self.RANGE_KEY = (self.RANGE_COL, self.RANGE_COL) self.KEYS = (self.HASH_COL, self.RANGE_COL) self.table_name = 'autotest_dynamo_db' self.dynamo_client = DynamoDbClient(config=self.TEST_CONFIG) def tearDown(self): clean_dynamo_table(self.table_name, self.KEYS) def test_put(self): row = { self.HASH_COL: 'cat', self.RANGE_COL: '123', 'some_bool': True, 'some_map': { 'a': 1, 'b': 'b1', 'c': { 'test': True } } } client = boto3.client('dynamodb') client.delete_item(TableName=self.table_name, Key={ self.HASH_COL: { 'S': str(row[self.HASH_COL]) }, self.RANGE_COL: { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) }, }) self.dynamo_client.put(row, self.table_name) result = client.scan( TableName=self.table_name, FilterExpression="hash_col = :hash_col AND range_col = :range_col", ExpressionAttributeValues={ ':hash_col': { 'S': row[self.HASH_COL] }, ':range_col': { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) } }) items = result['Items'] self.assertEqual(1, len(items)) expected = [{ 'hash_col': { 'S': 'cat' }, 'range_col': { 'N': '123' }, 'some_bool': { 'BOOL': True }, 'some_map': { 'M': { 'a': { 'N': '1' }, 'b': { 'S': 'b1' }, 'c': { 'M': { 'test': { 'BOOL': True } } } } } }] self.assertEqual(expected, items) def test_put__create(self): row = {self.HASH_COL: 'cat', self.RANGE_COL: '123'} self.dynamo_client.put(row, self.table_name) with self.assertRaises(self.dynamo_client.dynamo_client.exceptions. ConditionalCheckFailedException): self.dynamo_client.put(row, self.table_name, overwrite_existing=False) def test_update__updates(self): keys = {self.HASH_COL: 'cat', self.RANGE_COL: '123'} row = { self.HASH_COL: 'cat', self.RANGE_COL: '123', 'some_col': 'no', 'other_col': 'foo' } attributes_to_update = {'some_col': 'yes', 'new_col': 'yup'} self.dynamo_client.put(row, self.table_name) client = boto3.client('dynamodb') # First check that the row we are trying to update is PUT correctly. initial_row = client.get_item( Key={ self.HASH_COL: { 'S': row[self.HASH_COL] }, self.RANGE_COL: { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) } }, TableName=self.table_name, )['Item'] initial_row = self.dynamo_client.dynamo_to_dict(initial_row) self.assertIsNotNone(initial_row) self.assertEqual(initial_row['some_col'], 'no') self.assertEqual(initial_row['other_col'], 'foo') self.dynamo_client.update(keys, attributes_to_update, table_name=self.table_name) updated_row = client.get_item( Key={ self.HASH_COL: { 'S': row[self.HASH_COL] }, self.RANGE_COL: { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertIsNotNone(updated_row) self.assertEqual(updated_row['some_col'], 'yes'), "Updated field not really updated" self.assertEqual(updated_row['new_col'], 'yup'), "New field was not created" self.assertEqual( updated_row['other_col'], 'foo'), "This field should be preserved, update() damaged it" def test_update__increment(self): keys = {self.HASH_COL: 'cat', self.RANGE_COL: '123'} row = { self.HASH_COL: 'cat', self.RANGE_COL: '123', 'some_col': 'no', 'some_counter': 10 } attributes_to_increment = {'some_counter': '1'} self.dynamo_client.put(row, self.table_name) self.dynamo_client.update( keys, {}, attributes_to_increment=attributes_to_increment, table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ self.HASH_COL: { 'S': row[self.HASH_COL] }, self.RANGE_COL: { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertIsNotNone(updated_row) self.assertEqual(updated_row['some_counter'], 11) def test_update__increment_2(self): keys = {self.HASH_COL: 'cat', self.RANGE_COL: '123'} row = { self.HASH_COL: 'cat', self.RANGE_COL: '123', 'some_col': 'no', 'some_counter': 10 } attributes_to_increment = {'some_counter': 5} self.dynamo_client.put(row, self.table_name) self.dynamo_client.update( keys, {}, attributes_to_increment=attributes_to_increment, table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ self.HASH_COL: { 'S': row[self.HASH_COL] }, self.RANGE_COL: { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertIsNotNone(updated_row) self.assertEqual(updated_row['some_counter'], 15) def test_update__increment_no_default(self): keys = {self.HASH_COL: 'cat', self.RANGE_COL: '123'} row = {self.HASH_COL: 'cat', self.RANGE_COL: '123', 'some_col': 'no'} attributes_to_increment = {'some_counter': '3'} self.dynamo_client.put(row, self.table_name) self.dynamo_client.update( keys, {}, attributes_to_increment=attributes_to_increment, table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ self.HASH_COL: { 'S': row[self.HASH_COL] }, self.RANGE_COL: { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertIsNotNone(updated_row) self.assertEqual(updated_row['some_counter'], 3) def test_update__condition_expression(self): keys = {self.HASH_COL: 'slime', self.RANGE_COL: '41'} row = {self.HASH_COL: 'slime', self.RANGE_COL: '41', 'some_col': 'no'} self.dynamo_client.put(row, self.table_name) # Should fail because conditional expression does not match self.assertRaises(self.dynamo_client.dynamo_client.exceptions. ConditionalCheckFailedException, self.dynamo_client.update, keys, {}, attributes_to_increment={'some_counter': '3'}, condition_expression='some_col = yes', table_name=self.table_name) # Should pass self.dynamo_client.update( keys, {}, attributes_to_increment={'some_counter': '3'}, condition_expression='some_col = no', table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ self.HASH_COL: { 'S': row[self.HASH_COL] }, self.RANGE_COL: { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertEqual(updated_row['some_counter'], 3) def test_patch(self): keys = {self.HASH_COL: 'slime', self.RANGE_COL: '41'} row = {self.HASH_COL: 'slime', self.RANGE_COL: '41', 'some_col': 'no'} # Should fail because row doesn't exist self.assertRaises(self.dynamo_client.dynamo_client.exceptions. ConditionalCheckFailedException, self.dynamo_client.patch, keys, attributes_to_update={'some_col': 'yes'}, table_name=self.table_name) # Create the row self.dynamo_client.put(row, self.table_name) # Should pass because the row exists now self.dynamo_client.patch(keys, attributes_to_update={'some_col': 'yes'}, table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ self.HASH_COL: { 'S': row[self.HASH_COL] }, self.RANGE_COL: { self.RANGE_COL_TYPE: str(row[self.RANGE_COL]) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertEqual(updated_row['some_col'], 'yes') def test_get_by_query__primary_index(self): keys = {self.HASH_COL: 'cat', self.RANGE_COL: '123'} row = { self.HASH_COL: 'cat', self.RANGE_COL: 123, 'some_col': 'test', 'some_bool': True } self.dynamo_client.put(row, self.table_name) result = self.dynamo_client.get_by_query(keys=keys) self.assertEqual(len(result), 1) result = result[0] for key in row: self.assertEqual(row[key], result[key]) for key in result: self.assertEqual(row[key], result[key]) def test_get_by_query__primary_index__gets_multiple(self): row = {self.HASH_COL: 'cat', self.RANGE_COL: 123, 'some_col': 'test'} self.dynamo_client.put(row, self.table_name) row2 = { self.HASH_COL: 'cat', self.RANGE_COL: 1234, 'some_col': 'test2' } self.dynamo_client.put(row2, self.table_name) result = self.dynamo_client.get_by_query(keys={self.HASH_COL: 'cat'}) self.assertEqual(len(result), 2) result1 = [ x for x in result if x[self.RANGE_COL] == row[self.RANGE_COL] ][0] result2 = [ x for x in result if x[self.RANGE_COL] == row2[self.RANGE_COL] ][0] for key in row: self.assertEqual(row[key], result1[key]) for key in result1: self.assertEqual(row[key], result1[key]) for key in row2: self.assertEqual(row2[key], result2[key]) for key in result2: self.assertEqual(row2[key], result2[key]) def test_get_by_query__secondary_index(self): keys = {self.HASH_COL: 'cat', 'other_col': 'abc123'} row = { self.HASH_COL: 'cat', self.RANGE_COL: 123, 'other_col': 'abc123' } self.dynamo_client.put(row, self.table_name) result = self.dynamo_client.get_by_query(keys=keys, index_name='autotest_index') self.assertEqual(len(result), 1) result = result[0] for key in row: self.assertEqual(row[key], result[key]) for key in result: self.assertEqual(row[key], result[key]) def test_get_by_query__comparison(self): keys = {self.HASH_COL: 'cat', self.RANGE_COL: '300'} row1 = { self.HASH_COL: 'cat', self.RANGE_COL: 123, 'other_col': 'abc123' } row2 = { self.HASH_COL: 'cat', self.RANGE_COL: 456, 'other_col': 'abc123' } self.dynamo_client.put(row1, self.table_name) self.dynamo_client.put(row2, self.table_name) result = self.dynamo_client.get_by_query( keys=keys, comparisons={self.RANGE_COL: '<='}) self.assertEqual(len(result), 1) result = result[0] self.assertEqual(result, row1) def test_get_by_query__comparison_between(self): # Put sample data x = [ self.dynamo_client.put({ self.HASH_COL: 'cat', self.RANGE_COL: x }, self.table_name) for x in range(10) ] keys = { self.HASH_COL: 'cat', 'st_between_range_col': '3', 'en_between_range_col': '6' } result = self.dynamo_client.get_by_query( keys=keys, comparisons={self.RANGE_COL: 'between'}) # print(result) self.assertTrue(all(x[self.RANGE_COL] in range(3, 7) for x in result)) result = self.dynamo_client.get_by_query(keys=keys) # print(result) self.assertTrue(all(x[self.RANGE_COL] in range(3, 7) for x in result)), "Failed if unspecified comparison. " \ "Should be automatic for :st_between_..." def test_get_by_query__filter_expression(self): """ This _integration_ test runs multiple checks with same sample data for several comparators. Have a look at the manual if required: https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.OperatorsAndFunctions.html """ # Put sample data [ self.dynamo_client.put({ self.HASH_COL: 'cat', self.RANGE_COL: x }, self.table_name) for x in range(3) ] [ self.dynamo_client.put( { self.HASH_COL: 'cat', self.RANGE_COL: x, 'mark': 1 }, self.table_name) for x in range(3, 6) ] self.dynamo_client.put( { self.HASH_COL: 'cat', self.RANGE_COL: 6, 'mark': 0 }, self.table_name) self.dynamo_client.put( { self.HASH_COL: 'cat', self.RANGE_COL: 7, 'mark': 'a' }, self.table_name) # Condition by range_col will return five rows out of six: 0 - 4 # Filter expression neggs the first three rows because they don't have `mark = 1`. keys = {self.HASH_COL: 'cat', self.RANGE_COL: 4} result = self.dynamo_client.get_by_query( keys=keys, comparisons={self.RANGE_COL: '<='}, fetch_all_fields=True, filter_expression='mark = 1') # print(result) self.assertEqual(len(result), 2) self.assertEqual(result[0], { self.HASH_COL: 'cat', self.RANGE_COL: 3, 'mark': 1 }) self.assertEqual(result[1], { self.HASH_COL: 'cat', self.RANGE_COL: 4, 'mark': 1 }) # In the same test we check also some comparator _functions_. result = self.dynamo_client.get_by_query( keys=keys, comparisons={self.RANGE_COL: '<='}, fetch_all_fields=True, filter_expression='attribute_exists mark') # print(result) self.assertEqual(len(result), 2) self.assertEqual([x[self.RANGE_COL] for x in result], list(range(3, 5))) self.assertEqual(result[0], { self.HASH_COL: 'cat', self.RANGE_COL: 3, 'mark': 1 }) self.assertEqual(result[1], { self.HASH_COL: 'cat', self.RANGE_COL: 4, 'mark': 1 }) result = self.dynamo_client.get_by_query( keys=keys, comparisons={self.RANGE_COL: '<='}, fetch_all_fields=True, filter_expression='attribute_not_exists mark') # print(result) self.assertEqual(len(result), 3) self.assertEqual([x[self.RANGE_COL] for x in result], list(range(3))) def test_get_by_query__comparison_begins_with(self): self.table_name = 'autotest_config_component' # This table has a string range key self.HASH_KEY = ('env', 'S') self.RANGE_KEY = ('config_name', 'S') self.KEYS = ('env', 'config_name') config = { 'row_mapper': { 'env': 'S', 'config_name': 'S', 'config_value': 'S' }, 'required_fields': ['env', 'config_name', 'config_value'], 'table_name': 'autotest_config_component', 'hash_key': self.HASH_COL } self.dynamo_client = DynamoDbClient(config=config) row1 = { 'env': 'cat', 'config_name': 'testzing', 'config_value': 'abc123' } row2 = { 'env': 'cat', 'config_name': 'dont_get_this', 'config_value': 'abc123' } row3 = { 'env': 'cat', 'config_name': 'testzer', 'config_value': 'abc124' } self.dynamo_client.put(row1, self.table_name) self.dynamo_client.put(row2, self.table_name) self.dynamo_client.put(row3, self.table_name) keys = {'env': 'cat', 'config_name': 'testz'} result = self.dynamo_client.get_by_query( keys=keys, table_name=self.table_name, comparisons={'config_name': 'begins_with'}) self.assertEqual(len(result), 2) self.assertTrue(row1 in result) self.assertTrue(row3 in result) def test_get_by_query__max_items(self): # This function can also be used for some benchmarking, just change to bigger amounts manually. INITIAL_TASKS = 5 # Change to 500 to run benchmarking, and uncomment raise at the end of the test. for x in range(1000, 1000 + INITIAL_TASKS): row = {self.HASH_COL: f"key", self.RANGE_COL: x} self.dynamo_client.put(row, self.table_name) if INITIAL_TASKS > 10: time.sleep( 0.1 ) # Sleep a little to fit the Write Capacity (10 WCU) of autotest table. n = 3 st = time.perf_counter() result = self.dynamo_client.get_by_query({self.HASH_COL: 'key'}, table_name=self.table_name, max_items=n) bm = time.perf_counter() - st logging.info(f"Benchmark (n={n}): {bm}") self.assertEqual(len(result), n) self.assertLess(bm, 0.1) # Check unspecified limit. result = self.dynamo_client.get_by_query({self.HASH_COL: 'key'}, table_name=self.table_name) self.assertEqual(len(result), INITIAL_TASKS) def test_get_by_query__return_count(self): rows = [{ self.HASH_COL: 'cat1', self.RANGE_COL: 121, 'some_col': 'test1' }, { self.HASH_COL: 'cat1', self.RANGE_COL: 122, 'some_col': 'test2' }, { self.HASH_COL: 'cat1', self.RANGE_COL: 123, 'some_col': 'test3' }] for x in rows: self.dynamo_client.put(x, table_name=self.table_name) result = self.dynamo_client.get_by_query({self.HASH_COL: 'cat1'}, table_name=self.table_name, return_count=True) self.assertEqual(result, 3) def test_get_by_query__reverse(self): rows = [{ self.HASH_COL: 'cat1', self.RANGE_COL: 121, 'some_col': 'test1' }, { self.HASH_COL: 'cat1', self.RANGE_COL: 122, 'some_col': 'test2' }, { self.HASH_COL: 'cat1', self.RANGE_COL: 123, 'some_col': 'test3' }] for x in rows: self.dynamo_client.put(x, table_name=self.table_name) result = self.dynamo_client.get_by_query({self.HASH_COL: 'cat1'}, table_name=self.table_name, desc=True) self.assertEqual(result[0], rows[-1]) def test_get_by_scan__all(self): rows = [{ self.HASH_COL: 'cat1', self.RANGE_COL: 121, 'some_col': 'test1' }, { self.HASH_COL: 'cat2', self.RANGE_COL: 122, 'some_col': 'test2' }, { self.HASH_COL: 'cat3', self.RANGE_COL: 123, 'some_col': 'test3' }] for x in rows: self.dynamo_client.put(x, self.table_name) result = self.dynamo_client.get_by_scan() self.assertEqual(len(result), 3) for r in rows: assert r in result, f"row not in result from dynamo scan: {r}" def test_get_by_scan__with_filter(self): rows = [ { self.HASH_COL: 'cat1', self.RANGE_COL: 121, 'some_col': 'test1' }, { self.HASH_COL: 'cat1', self.RANGE_COL: 122, 'some_col': 'test2' }, { self.HASH_COL: 'cat2', self.RANGE_COL: 122, 'some_col': 'test2' }, ] for x in rows: self.dynamo_client.put(x, self.table_name) filter = {'some_col': 'test2'} result = self.dynamo_client.get_by_scan(attrs=filter) self.assertEqual(len(result), 2) for r in rows[1:]: assert r in result, f"row not in result from dynamo scan: {r}" def test_batch_get_items(self): rows = [ { self.HASH_COL: 'cat1', self.RANGE_COL: 121, 'some_col': 'test1' }, { self.HASH_COL: 'cat1', self.RANGE_COL: 122, 'some_col': 'test2' }, { self.HASH_COL: 'cat2', self.RANGE_COL: 122, 'some_col': 'test2' }, ] for x in rows: self.dynamo_client.put(x, self.table_name) keys_list_query = [ { self.HASH_COL: 'cat1', self.RANGE_COL: 121 }, { self.HASH_COL: 'doesnt_exist', self.RANGE_COL: 40 }, { self.HASH_COL: 'cat2', self.RANGE_COL: 122 }, ] result = self.dynamo_client.batch_get_items_one_table(keys_list_query) self.assertEqual(len(result), 2) self.assertIn(rows[0], result) self.assertIn(rows[2], result) def test_delete(self): self.dynamo_client.put({self.HASH_COL: 'cat1', self.RANGE_COL: 123}) self.dynamo_client.put({self.HASH_COL: 'cat2', self.RANGE_COL: 234}) self.dynamo_client.delete(keys={ self.HASH_COL: 'cat1', self.RANGE_COL: '123' }) items = self.dynamo_client.get_by_scan() self.assertEqual(len(items), 1) self.assertEqual(items[0], { self.HASH_COL: 'cat2', self.RANGE_COL: 234 }) def test_get_table_keys(self): result1 = self.dynamo_client.get_table_keys() self.assertEqual(result1, ('hash_col', 'range_col')) result2 = self.dynamo_client.get_table_keys(table_name=self.table_name) self.assertEqual(result2, ('hash_col', 'range_col')) def test_get_table_indexes(self): indexes = self.dynamo_client.get_table_indexes() expected = { 'autotest_index': { 'projection_type': 'ALL', 'hash_key': 'hash_col', 'range_key': 'other_col', 'provisioned_throughput': { 'write_capacity': 1, 'read_capacity': 1 } } } self.assertDictEqual(expected, indexes) def test_batch_get_items_one_table(self): # If you want to stress test batch_get_items_one_table, use bigger numbers num_of_items = 5 query_from = 2 query_till = 4 expected_items = query_till - query_from # Write items operations = [] query_keys = [] for i in range(num_of_items): item = {self.HASH_COL: f'cat{i%2}', self.RANGE_COL: i} operations.append( {'Put': self.dynamo_client.build_put_query(item)}) query_keys.append(item) for operations_chunk in chunks(operations, 10): self.dynamo_client.dynamo_client.transact_write_items( TransactItems=operations_chunk) time.sleep(1) # cause the table has 10 write/sec capacity # Batch get items results = self.dynamo_client.batch_get_items_one_table( keys_list=query_keys[query_from:query_till]) self.assertEqual(expected_items, len(results))
class dynamodb_client_IntegrationTestCase(unittest.TestCase): TEST_CONFIG = { 'row_mapper': { 'lambda_name': 'S', 'invocation_id': 'S', 'en_time': 'N', 'hash_col': 'S', 'range_col': 'N', 'other_col': 'S', 'new_col': 'S', 'some_col': 'S', 'some_counter': 'N' }, 'required_fields': ['lambda_name'], 'table_name': 'autotest_dynamo_db' } @classmethod def setUpClass(cls): clean_dynamo_table() def setUp(self): self.HASH_KEY = ('hash_col', 'S') self.RANGE_KEY = ('range_col', 'N') self.KEYS = ('hash_col', 'range_col') self.table_name = 'autotest_dynamo_db' self.dynamo_client = DynamoDbClient(config=self.TEST_CONFIG) def tearDown(self): clean_dynamo_table(self.table_name, self.KEYS) def test_put(self): row = {'hash_col': 'cat', 'range_col': '123'} client = boto3.client('dynamodb') client.delete_item(TableName=self.table_name, Key={ 'hash_col': { 'S': str(row['hash_col']) }, 'range_col': { 'N': str(row['range_col']) }, }) self.dynamo_client.put(row, self.table_name) result = client.scan( TableName=self.table_name, FilterExpression="hash_col = :hash_col AND range_col = :range_col", ExpressionAttributeValues={ ':hash_col': { 'S': row['hash_col'] }, ':range_col': { 'N': str(row['range_col']) } }) items = result['Items'] self.assertTrue(len(items) > 0) def test_update__updates(self): keys = {'hash_col': 'cat', 'range_col': '123'} row = { 'hash_col': 'cat', 'range_col': '123', 'some_col': 'no', 'other_col': 'foo' } attributes_to_update = {'some_col': 'yes', 'new_col': 'yup'} self.dynamo_client.put(row, self.table_name) client = boto3.client('dynamodb') # First check that the row we are trying to update is PUT correctly. initial_row = client.get_item( Key={ 'hash_col': { 'S': row['hash_col'] }, 'range_col': { 'N': str(row['range_col']) } }, TableName=self.table_name, )['Item'] initial_row = self.dynamo_client.dynamo_to_dict(initial_row) self.assertIsNotNone(initial_row) self.assertEqual(initial_row['some_col'], 'no') self.assertEqual(initial_row['other_col'], 'foo') self.dynamo_client.update(keys, attributes_to_update, table_name=self.table_name) updated_row = client.get_item( Key={ 'hash_col': { 'S': row['hash_col'] }, 'range_col': { 'N': str(row['range_col']) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertIsNotNone(updated_row) self.assertEqual(updated_row['some_col'], 'yes'), "Updated field not really updated" self.assertEqual(updated_row['new_col'], 'yup'), "New field was not created" self.assertEqual( updated_row['other_col'], 'foo'), "This field should be preserved, update() damaged it" def test_update__increment(self): keys = {'hash_col': 'cat', 'range_col': '123'} row = { 'hash_col': 'cat', 'range_col': '123', 'some_col': 'no', 'some_counter': 10 } attributes_to_increment = {'some_counter': '1'} self.dynamo_client.put(row, self.table_name) self.dynamo_client.update( keys, {}, attributes_to_increment=attributes_to_increment, table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ 'hash_col': { 'S': row['hash_col'] }, 'range_col': { 'N': str(row['range_col']) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertIsNotNone(updated_row) self.assertEqual(updated_row['some_counter'], 11) def test_update__increment_2(self): keys = {'hash_col': 'cat', 'range_col': '123'} row = { 'hash_col': 'cat', 'range_col': '123', 'some_col': 'no', 'some_counter': 10 } attributes_to_increment = {'some_counter': 5} self.dynamo_client.put(row, self.table_name) self.dynamo_client.update( keys, {}, attributes_to_increment=attributes_to_increment, table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ 'hash_col': { 'S': row['hash_col'] }, 'range_col': { 'N': str(row['range_col']) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertIsNotNone(updated_row) self.assertEqual(updated_row['some_counter'], 15) def test_update__increment_no_default(self): keys = {'hash_col': 'cat', 'range_col': '123'} row = {'hash_col': 'cat', 'range_col': '123', 'some_col': 'no'} attributes_to_increment = {'some_counter': '3'} self.dynamo_client.put(row, self.table_name) self.dynamo_client.update( keys, {}, attributes_to_increment=attributes_to_increment, table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ 'hash_col': { 'S': row['hash_col'] }, 'range_col': { 'N': str(row['range_col']) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertIsNotNone(updated_row) self.assertEqual(updated_row['some_counter'], 3) def test_update__condition_expression(self): keys = {'hash_col': 'slime', 'range_col': '41'} row = {'hash_col': 'slime', 'range_col': '41', 'some_col': 'no'} self.dynamo_client.put(row, self.table_name) # Should fail because conditional expression does not match self.assertRaises(self.dynamo_client.dynamo_client.exceptions. ConditionalCheckFailedException, self.dynamo_client.update, keys, {}, attributes_to_increment={'some_counter': '3'}, condition_expression='some_col = yes', table_name=self.table_name) # Should pass self.dynamo_client.update( keys, {}, attributes_to_increment={'some_counter': '3'}, condition_expression='some_col = no', table_name=self.table_name) client = boto3.client('dynamodb') updated_row = client.get_item( Key={ 'hash_col': { 'S': row['hash_col'] }, 'range_col': { 'N': str(row['range_col']) } }, TableName=self.table_name, )['Item'] updated_row = self.dynamo_client.dynamo_to_dict(updated_row) self.assertEqual(updated_row['some_counter'], 3) def test_get_by_query__primary_index(self): keys = {'hash_col': 'cat', 'range_col': '123'} row = {'hash_col': 'cat', 'range_col': 123, 'some_col': 'test'} self.dynamo_client.put(row, self.table_name) result = self.dynamo_client.get_by_query(keys=keys) self.assertEqual(len(result), 1) result = result[0] for key in row: self.assertEqual(row[key], result[key]) for key in result: self.assertEqual(row[key], result[key]) def test_get_by_query__primary_index__gets_multiple(self): row = {'hash_col': 'cat', 'range_col': 123, 'some_col': 'test'} self.dynamo_client.put(row, self.table_name) row2 = {'hash_col': 'cat', 'range_col': 1234, 'some_col': 'test2'} self.dynamo_client.put(row2, self.table_name) result = self.dynamo_client.get_by_query(keys={'hash_col': 'cat'}) self.assertEqual(len(result), 2) result1 = [x for x in result if x['range_col'] == row['range_col']][0] result2 = [x for x in result if x['range_col'] == row2['range_col']][0] for key in row: self.assertEqual(row[key], result1[key]) for key in result1: self.assertEqual(row[key], result1[key]) for key in row2: self.assertEqual(row2[key], result2[key]) for key in result2: self.assertEqual(row2[key], result2[key]) def test_get_by_query__secondary_index(self): keys = {'hash_col': 'cat', 'other_col': 'abc123'} row = {'hash_col': 'cat', 'range_col': 123, 'other_col': 'abc123'} self.dynamo_client.put(row, self.table_name) result = self.dynamo_client.get_by_query(keys=keys, index_name='autotest_index') self.assertEqual(len(result), 1) result = result[0] for key in row: self.assertEqual(row[key], result[key]) for key in result: self.assertEqual(row[key], result[key]) def test_get_by_query__comparison(self): keys = {'hash_col': 'cat', 'range_col': '300'} row1 = {'hash_col': 'cat', 'range_col': 123, 'other_col': 'abc123'} row2 = {'hash_col': 'cat', 'range_col': 456, 'other_col': 'abc123'} self.dynamo_client.put(row1, self.table_name) self.dynamo_client.put(row2, self.table_name) result = self.dynamo_client.get_by_query( keys=keys, comparisons={'range_col': '<='}) self.assertEqual(len(result), 1) result = result[0] self.assertEqual(result, row1) def test_get_by_query__comparison_between(self): # Put sample data x = [ self.dynamo_client.put({ 'hash_col': 'cat', 'range_col': x }, self.table_name) for x in range(10) ] keys = { 'hash_col': 'cat', 'st_between_range_col': '3', 'en_between_range_col': '6' } result = self.dynamo_client.get_by_query( keys=keys, comparisons={'range_col': 'between'}) # print(result) self.assertTrue(all(x['range_col'] in range(3, 7) for x in result)) result = self.dynamo_client.get_by_query(keys=keys) # print(result) self.assertTrue(all(x['range_col'] in range(3, 7) for x in result)), "Failed if unspecified comparison. " \ "Should be automatic for :st_between_..." def test_get_by_query__filter_expression(self): """ This _integration_ test runs multiple checks with same sample data for several comparators. Have a look at the manual if required: https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Expressions.OperatorsAndFunctions.html """ # Put sample data [ self.dynamo_client.put({ 'hash_col': 'cat', 'range_col': x }, self.table_name) for x in range(3) ] [ self.dynamo_client.put( { 'hash_col': 'cat', 'range_col': x, 'mark': 1 }, self.table_name) for x in range(3, 6) ] self.dynamo_client.put({ 'hash_col': 'cat', 'range_col': 6, 'mark': 0 }, self.table_name) self.dynamo_client.put({ 'hash_col': 'cat', 'range_col': 7, 'mark': 'a' }, self.table_name) # Condition by range_col will return five rows out of six: 0 - 4 # Filter expression neggs the first three rows because they don't have `mark = 1`. keys = {'hash_col': 'cat', 'range_col': 4} result = self.dynamo_client.get_by_query( keys=keys, comparisons={'range_col': '<='}, strict=False, filter_expression='mark = 1') # print(result) self.assertEqual(len(result), 2) self.assertEqual(result[0], { 'hash_col': 'cat', 'range_col': 3, 'mark': 1 }) self.assertEqual(result[1], { 'hash_col': 'cat', 'range_col': 4, 'mark': 1 }) # In the same test we check also some comparator _functions_. result = self.dynamo_client.get_by_query( keys=keys, comparisons={'range_col': '<='}, strict=False, filter_expression='attribute_exists mark') # print(result) self.assertEqual(len(result), 2) self.assertEqual([x['range_col'] for x in result], list(range(3, 5))) self.assertEqual(result[0], { 'hash_col': 'cat', 'range_col': 3, 'mark': 1 }) self.assertEqual(result[1], { 'hash_col': 'cat', 'range_col': 4, 'mark': 1 }) result = self.dynamo_client.get_by_query( keys=keys, comparisons={'range_col': '<='}, strict=False, filter_expression='attribute_not_exists mark') # print(result) self.assertEqual(len(result), 3) self.assertEqual([x['range_col'] for x in result], list(range(3))) def test_get_by_query__comparison_begins_with(self): self.table_name = 'autotest_config_component' # This table has a string range key self.HASH_KEY = ('env', 'S') self.RANGE_KEY = ('config_name', 'S') self.KEYS = ('env', 'config_name') config = { 'row_mapper': { 'env': 'S', 'config_name': 'S', 'config_value': 'S' }, 'required_fields': ['env', 'config_name', 'config_value'], 'table_name': 'autotest_config_component' } self.dynamo_client = DynamoDbClient(config=config) row1 = { 'env': 'cat', 'config_name': 'testzing', 'config_value': 'abc123' } row2 = { 'env': 'cat', 'config_name': 'dont_get_this', 'config_value': 'abc123' } row3 = { 'env': 'cat', 'config_name': 'testzer', 'config_value': 'abc124' } self.dynamo_client.put(row1, self.table_name) self.dynamo_client.put(row2, self.table_name) self.dynamo_client.put(row3, self.table_name) keys = {'env': 'cat', 'config_name': 'testz'} result = self.dynamo_client.get_by_query( keys=keys, table_name=self.table_name, comparisons={'config_name': 'begins_with'}) self.assertEqual(len(result), 2) self.assertTrue(row1 in result) self.assertTrue(row3 in result) def test_get_by_query__max_items(self): # This function can also be used for some benchmarking, just change to bigger amounts manually. INITIAL_TASKS = 5 # Change to 500 to run benchmarking, and uncomment raise at the end of the test. for x in range(1000, 1000 + INITIAL_TASKS): row = {'hash_col': f"key", 'range_col': x} self.dynamo_client.put(row, self.table_name) if INITIAL_TASKS > 10: time.sleep( 0.1 ) # Sleep a little to fit the Write Capacity (10 WCU) of autotest table. st = time.perf_counter() result = self.dynamo_client.get_by_query({'hash_col': 'key'}, table_name=self.table_name, max_items=3) bm = time.perf_counter() - st print(f"Benchmark: {bm}") self.assertEqual(len(result), 3) self.assertLess(bm, 0.1) # Check unspecified limit. result = self.dynamo_client.get_by_query({'hash_col': 'key'}, table_name=self.table_name) self.assertEqual(len(result), INITIAL_TASKS) # Benchmarking if INITIAL_TASKS >= 500: st = time.perf_counter() result = self.dynamo_client.get_by_query( {'hash_col': 'key'}, table_name=self.table_name, max_items=499) bm = time.perf_counter() - st print(f"Benchmark: {bm}") self.assertLess(bm, 0.1) self.assertEqual(len(result), 499) # Uncomment this see benchmark # self.assertEqual(1, 2) def test_get_by_query__return_count(self): rows = [{ 'hash_col': 'cat1', 'range_col': 121, 'some_col': 'test1' }, { 'hash_col': 'cat1', 'range_col': 122, 'some_col': 'test2' }, { 'hash_col': 'cat1', 'range_col': 123, 'some_col': 'test3' }] for x in rows: self.dynamo_client.put(x, table_name=self.table_name) result = self.dynamo_client.get_by_query({'hash_col': 'cat1'}, table_name=self.table_name, return_count=True) self.assertEqual(result, 3) def test_get_by_query__reverse(self): rows = [{ 'hash_col': 'cat1', 'range_col': 121, 'some_col': 'test1' }, { 'hash_col': 'cat1', 'range_col': 122, 'some_col': 'test2' }, { 'hash_col': 'cat1', 'range_col': 123, 'some_col': 'test3' }] for x in rows: self.dynamo_client.put(x, table_name=self.table_name) result = self.dynamo_client.get_by_query({'hash_col': 'cat1'}, table_name=self.table_name, desc=True) self.assertEqual(result[0], rows[-1]) def test_get_by_scan__all(self): rows = [{ 'hash_col': 'cat1', 'range_col': 121, 'some_col': 'test1' }, { 'hash_col': 'cat2', 'range_col': 122, 'some_col': 'test2' }, { 'hash_col': 'cat3', 'range_col': 123, 'some_col': 'test3' }] for x in rows: self.dynamo_client.put(x, self.table_name) result = self.dynamo_client.get_by_scan() self.assertEqual(len(result), 3) for r in rows: assert r in result, f"row not in result from dynamo scan: {r}" def test_get_by_scan__with_filter(self): rows = [ { 'hash_col': 'cat1', 'range_col': 121, 'some_col': 'test1' }, { 'hash_col': 'cat1', 'range_col': 122, 'some_col': 'test2' }, { 'hash_col': 'cat2', 'range_col': 122, 'some_col': 'test2' }, ] for x in rows: self.dynamo_client.put(x, self.table_name) filter = {'some_col': 'test2'} result = self.dynamo_client.get_by_scan(attrs=filter) self.assertEqual(len(result), 2) for r in rows[1:]: assert r in result, f"row not in result from dynamo scan: {r}" def test_batch_get_items(self): rows = [ { 'hash_col': 'cat1', 'range_col': 121, 'some_col': 'test1' }, { 'hash_col': 'cat1', 'range_col': 122, 'some_col': 'test2' }, { 'hash_col': 'cat2', 'range_col': 122, 'some_col': 'test2' }, ] for x in rows: self.dynamo_client.put(x, self.table_name) keys_list_query = [ { 'hash_col': 'cat1', 'range_col': 121 }, { 'hash_col': 'doesnt_exist', 'range_col': 40 }, { 'hash_col': 'cat2', 'range_col': 122 }, ] result = self.dynamo_client.batch_get_items_one_table(keys_list_query) self.assertEqual(len(result), 2) self.assertIn(rows[0], result) self.assertIn(rows[2], result) def test_delete(self): self.dynamo_client.put({'hash_col': 'cat1', 'range_col': 123}) self.dynamo_client.put({'hash_col': 'cat2', 'range_col': 234}) self.dynamo_client.delete(keys={ 'hash_col': 'cat1', 'range_col': '123' }) items = self.dynamo_client.get_by_scan() self.assertEqual(len(items), 1) self.assertEqual(items[0], {'hash_col': 'cat2', 'range_col': 234})