''' Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved. SPDX-License-Identifier: MIT-0 ''' import os import boto3 def lambda_handler(event, context): operation = event['operation'] client = boto3.client('lambda') config = client.get_function_configuration( FunctionName=os.environ['WORKLOAD_ARN'] ) env_vars = config['Environment'] if operation == 'start': print('Starting workload') env_vars['Variables']['IS_ACTIVE'] = 'True' client.update_function_configuration( FunctionName=os.environ['WORKLOAD_ARN'], Environment=env_vars ) client.invoke( FunctionName=os.environ['WORKLOAD_ARN'], InvocationType='Event' ) else: print('Stopping workload') env_vars['Variables']['IS_ACTIVE'] = 'False' client.update_function_configuration( FunctionName=os.environ['WORKLOAD_ARN'], Environment=env_vars )