def lambda_context():
context = LambdaContext()
context.aws_request_id = 'test_aws_request_id'
context.function_name = 'emr_create_function_name'
context.log_group_name = 'log_group_name'
context.log_stream_name = 'log_stream_name'
return context
def lambda_ctx() -> LambdaContext:
lambda_cognito_identity = LambdaCognitoIdentity()
lambda_cognito_identity.cognito_identity_id = "cognito_identity_id"
lambda_cognito_identity.cognito_identity_pool_id = "cognito_identity_pool_id"
lambda_client_context_mobile_client = LambdaClientContextMobileClient()
lambda_client_context_mobile_client.installation_id = "installation_id"
lambda_client_context_mobile_client.app_title = "app_title"
lambda_client_context_mobile_client.app_version_name = "app_version_name"
lambda_client_context_mobile_client.app_version_code = "app_version_code"
lambda_client_context_mobile_client.app_package_name = "app_package_name"
lambda_client_context = LambdaClientContext()
lambda_client_context.client = lambda_client_context_mobile_client
lambda_client_context.custom = {"custom": True}
lambda_client_context.env = {"env": "test"}
lambda_context = LambdaContext()
lambda_context.function_name = "function_name"
lambda_context.function_version = "function_version"
lambda_context.invoked_function_arn = "invoked_function_arn"
lambda_context.memory_limit_in_mb = 1234
lambda_context.aws_request_id = "aws_request_id"
lambda_context.log_group_name = "log_group_name"
lambda_context.log_stream_name = "log_stream_name"
lambda_context.identity = lambda_cognito_identity
lambda_context.client_context = lambda_client_context
return lambda_context
def mock_context():
context = LambdaContext()
context.function_name = 'test'
context.function_version = 'test'
context.invoked_function_arn = 'test'
context.memory_limit_in_mb = 'test'
context.aws_request_id = 'test'
context.log_group_name = 'test'
context.log_stream_name = 'test'
return context
def test_all_private_ips_and_shards_matching_clickhouse_server_added_to_zookeeper(
self,
mock_get_ec2_client,
mock_get_zookeeper_client,
mock_get_clickhouse_cluster_definition,
):
ec2_client = MagicMock()
mock_get_ec2_client.return_value = ec2_client
shard_config = {
"shard_1": ["172.26.39.30", "172.26.99.237"],
"shard_2": ["172.26.32.16", "172.26.97.29"],
}
mock_get_clickhouse_cluster_definition.return_value = shard_config
zookeeper = MagicMock()
mock_get_zookeeper_client.return_value = zookeeper
lambda_context = LambdaContext()
lambda_context.function_name = "lambda_handler"
lambda_context.aws_request_id = "abc-123"
result = lambda_handler({}, lambda_context)
remote_servers_xml = b"""<hmrc_data_cluster>
<shard>
<internal_replication>true</internal_replication>
<replica>
<default_database>graphite</default_database>
<host>172.26.39.30</host>
<port>9000</port>
</replica>
<replica>
<default_database>graphite</default_database>
<host>172.26.99.237</host>
<port>9000</port>
</replica>
</shard>
<shard>
<internal_replication>true</internal_replication>
<replica>
<default_database>graphite</default_database>
<host>172.26.32.16</host>
<port>9000</port>
</replica>
<replica>
<default_database>graphite</default_database>
<host>172.26.97.29</host>
<port>9000</port>
</replica>
</shard>
</hmrc_data_cluster>"""
zookeeper.set.assert_any_call("clickhouse.config.remote_servers",
remote_servers_xml)
self.assertEqual(result, {"cluster_definition": shard_config})
def test_should_ensure_zookeeper_path_exists_for_remote_servers(
self,
mock_get_ec2_client,
mock_get_zookeeper_client,
mock_get_clickhouse_cluster_definition,
):
zookeeper = MagicMock()
mock_get_zookeeper_client.return_value = zookeeper
lambda_context = LambdaContext()
lambda_context.function_name = "lambda_handler"
lambda_context.aws_request_id = "abc-123"
lambda_handler({}, lambda_context)
zookeeper.ensure_path.assert_any_call(
"clickhouse.config.remote_servers")
def test_that_the_lambda_handler_succeeds_with_context(sns_event):
lambda_context = LambdaContext()
lambda_context.function_name = 'lambda_handler'
lambda_context.aws_request_id = 'abc-123'
ecs_client = boto3.client('ecs')
ecs_cluster_name = get_ecs_cluster_name()
riemann_consumer_ecs_service_name = get_riemann_consumer_ecs_service_name()
ecs_client.create_cluster(clusterName=ecs_cluster_name)
ecs_client.create_service(cluster=ecs_cluster_name,
serviceName=riemann_consumer_ecs_service_name)
response = lambda_handler(event=sns_event, context=lambda_context)
assert response['success'] is True
assert response['serviceName'] == riemann_consumer_ecs_service_name
assert response['status'] == "ACTIVE"