def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if 'meta_handler_config' in self.config:
self.meta_handler = MetaHandler(
custom_config=self.config['meta_handler_config'])
def setUp(self):
self.patcher = patch("sosw.app.get_config")
self.get_config_patch = self.patcher.start()
self.config = deepcopy(TEST_META_HANDLER_CONFIG)
with patch('boto3.client'):
global_vars.lambda_context = self.TEST_CONTEXT
self.manager = MetaHandler(custom_config=self.config)
self.manager.dynamo_db_client = MagicMock(
spec=dynamo_db.DynamoDbClient)
self.manager.lambda_context = TEST_META_HANDLER_LAMBDA_CONTEXT
self.expected_values = {
'task_id': TEST_META_HANDLER_POST_ARGS['task_id'],
'action': self.manager._ma(TEST_META_HANDLER_POST_ARGS['action']),
'author': 'author',
'invocation_id': 'invocation_id',
'log_stream_name': 'log_stream_name'
}
class Worker(Processor):
"""
We recommend that you inherit your core Processor from this class in Lambdas that are orchestrated by `sosw`.
The ``__call__`` method is supposed to accept the ``event`` of the Lambda invocation.
This is a dictionary with the payload received in the lambda_handler during invocation.
Worker has all the common methods of :ref:`Processor` and tries to mark task as completed if received
``task_id`` in the ``event``. Worker create a payload with ``stats`` and ``result`` if exist and invoke worker
assistant lambda.
Worker class can optionally record ``'completed'`` and ``'failed'`` events to the DynamoDB tasks meta data table.
In order to enable this feature, you have to provide ``'meta_handler_config'`` in your custom_config.
You also need to grant write permissions for this table to your Lambda.
You can find more information about the configuration in the :ref:`MetaHandler<meta_handler>` chapter.
"""
DEFAULT_CONFIG = {
'init_clients': ['lambda'],
'sosw_worker_assistant_lambda': 'sosw_worker_assistant'
}
lambda_client = None
meta_handler: MetaHandler = None
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if 'meta_handler_config' in self.config:
self.meta_handler = MetaHandler(
custom_config=self.config['meta_handler_config'])
def __call__(self, event: Dict):
"""
You can either call super() at the end of your child function or completely overwrite this function.
"""
# Mark the task as completed in DynamoDB if the event had task_id.
try:
if event.get('task_id'):
self.mark_task_as_completed(event['task_id'])
except Exception:
logger.exception(
f"Failed to call WorkerAssistant for event {event}")
pass
super().__call__(event)
def mark_task_as_completed(self, task_id: str):
""" Call worker assistant lambda and tell it to close task """
if not self.lambda_client:
self.register_clients(['lambda'])
worker_assistant_lambda_name = self.config.get(
'sosw_worker_assistant_lambda', 'sosw_worker_assistant')
payload = {
'action': 'mark_task_as_completed',
'task_id': task_id,
}
if self.stats:
payload.update({'stats': self.stats})
if self.result:
payload.update({'result': self.result})
payload = json.dumps(payload)
lambda_response = self.lambda_client.invoke(
FunctionName=worker_assistant_lambda_name,
InvocationType='Event',
Payload=payload)
if self.meta_handler:
self.meta_handler.post(task_id=task_id, action='completed')
logger.debug(f"mark_task_as_completed response: {lambda_response}")
def mark_task_as_failed(self, task_id: str):
""" Call worker assistant lambda and tell it to update task info """
if not self.lambda_client:
self.register_clients(['lambda'])
worker_assistant_lambda_name = self.config.get(
'sosw_worker_assistant_lambda', 'sosw_worker_assistant')
payload = {
'action': 'mark_task_as_failed',
'task_id': task_id,
}
if self.stats:
payload.update({'stats': self.stats})
if self.result:
payload.update({'result': self.result})
payload = json.dumps(payload)
lambda_response = self.lambda_client.invoke(
FunctionName=worker_assistant_lambda_name,
InvocationType='Event',
Payload=payload)
if self.meta_handler:
self.meta_handler.post(task_id=task_id, action='failed')
logger.debug(f"mark_task_as_failed response: {lambda_response}")
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.meta_handler = MetaHandler(
custom_config=self.config.get('meta_handler_config'))
class meta_handler_UnitTestCase(unittest.TestCase):
TEST_CONTEXT = attrdict.AttrDict(
{v: k
for k, v in MetaHandler.CONTEXT_FIELDS_MAPPINGS.items()})
def setUp(self):
self.patcher = patch("sosw.app.get_config")
self.get_config_patch = self.patcher.start()
self.config = deepcopy(TEST_META_HANDLER_CONFIG)
with patch('boto3.client'):
global_vars.lambda_context = self.TEST_CONTEXT
self.manager = MetaHandler(custom_config=self.config)
self.manager.dynamo_db_client = MagicMock(
spec=dynamo_db.DynamoDbClient)
self.manager.lambda_context = TEST_META_HANDLER_LAMBDA_CONTEXT
self.expected_values = {
'task_id': TEST_META_HANDLER_POST_ARGS['task_id'],
'action': self.manager._ma(TEST_META_HANDLER_POST_ARGS['action']),
'author': 'author',
'invocation_id': 'invocation_id',
'log_stream_name': 'log_stream_name'
}
def tearDown(self):
self.patcher.stop()
def test_post__raise_kwargs_intersection_with_lambda_context(self):
"""
Method should raise in case attempt to overwrite lambda_context fields from kwargs
"""
kwargs = {'invocation_id': 'test_invocation_id'}
self.assertRaises(AssertionError, self.manager.post,
**TEST_META_HANDLER_POST_ARGS, **kwargs)
def test_post__check_create_call(self):
"""
Check that method will try to create correct row
"""
self.manager.post(**TEST_META_HANDLER_POST_ARGS)
self.manager.dynamo_db_client.create.assert_called_once()
args, kwargs = self.manager.dynamo_db_client.create.call_args
for k, v in self.expected_values.items():
self.assertEqual(v, kwargs['row'][k])
def test_post__check_create_call_with_args_from_kwargs(self):
"""
Check that method will try to create correct row with kwargs
"""
kwargs = {
'result': 'test_results',
'stats': 'test_stats',
'health_status': 'healthy'
}
expected_row = deepcopy(self.expected_values)
expected_row.update(kwargs)
self.manager.post(**TEST_META_HANDLER_POST_ARGS, **kwargs)
self.manager.dynamo_db_client.create.assert_called_once()
args, kwargs = self.manager.dynamo_db_client.create.call_args
for k, v in self.expected_values.items():
self.assertEqual(v, kwargs['row'][k])
self.assertEqual(type(self.manager._ma(123)), str)