diff --git a/lib/aws_lambda_ric/lambda_context.rb b/lib/aws_lambda_ric/lambda_context.rb index 52758a8..259ba16 100644 --- a/lib/aws_lambda_ric/lambda_context.rb +++ b/lib/aws_lambda_ric/lambda_context.rb @@ -3,7 +3,7 @@ class LambdaContext attr_reader :aws_request_id, :invoked_function_arn, :log_group_name, :log_stream_name, :function_name, :memory_limit_in_mb, :function_version, - :identity, :client_context, :deadline_ms + :identity, :tenant_id, :client_context, :deadline_ms def initialize(request) @clock_diff = Process.clock_gettime(Process::CLOCK_REALTIME, :millisecond) - Process.clock_gettime(Process::CLOCK_MONOTONIC, :millisecond) @@ -16,6 +16,7 @@ def initialize(request) @memory_limit_in_mb = ENV['AWS_LAMBDA_FUNCTION_MEMORY_SIZE'] @function_version = ENV['AWS_LAMBDA_FUNCTION_VERSION'] @identity = JSON.parse(request['Lambda-Runtime-Cognito-Identity']) unless request['Lambda-Runtime-Cognito-Identity'].to_s.empty? + @tenant_id = request['Lambda-Runtime-Aws-Tenant-Id'] unless request['Lambda-Runtime-Aws-Tenant-Id'].to_s.empty? @client_context = JSON.parse(request['Lambda-Runtime-Client-Context']) unless request['Lambda-Runtime-Client-Context'].to_s.empty? end diff --git a/test/unit/lambda_server_test.rb b/test/unit/lambda_server_test.rb index 40d09db..fb2e248 100644 --- a/test/unit/lambda_server_test.rb +++ b/test/unit/lambda_server_test.rb @@ -53,4 +53,62 @@ def test_post_invocation_error_with_too_large_xray_cause assert_mock post_mock end + + def mock_next_invocation_response() + mock_response = Net::HTTPSuccess.new(1.0, '200', 'OK') + mock_response['Lambda-Runtime-Aws-Request-Id'] = @request_id + mock_response + end + + def mock_next_invocation_request(mock_response) + get_mock = Minitest::Mock.new + get_mock.expect(:read_timeout=, nil, [RapidClient::LONG_TIMEOUT_MS]) + get_mock.expect(:start, mock_response) do |&block| + block.call(get_mock) + end + get_mock.expect(:get, mock_response, ['/2018-06-01/runtime/invocation/next', {'User-Agent' => @mock_user_agent}]) + get_mock + end + + def assert_next_invocation(get_mock, expected_tenant_id) + Net::HTTP.stub(:new, get_mock, ['127.0.0.1', 9001]) do + request_id, response = @under_test.next_invocation + assert_equal @request_id, request_id + assert_equal expected_tenant_id, response['Lambda-Runtime-Aws-Tenant-Id'] + end + end + + def test_next_invocation_without_tenant_id_header + mock_response = mock_next_invocation_response() + get_mock = mock_next_invocation_request(mock_response) + assert_next_invocation(get_mock, nil) + assert_mock get_mock + end + + def test_next_invocation_with_tenant_id_header + mock_response = mock_next_invocation_response() + mock_response['Lambda-Runtime-Aws-Tenant-Id'] = 'blue' + + get_mock = mock_next_invocation_request(mock_response) + assert_next_invocation(get_mock, 'blue') + assert_mock get_mock + end + + def test_next_invocation_with_empty_tenant_id_header + mock_response = mock_next_invocation_response() + mock_response['Lambda-Runtime-Aws-Tenant-Id'] = '' + + get_mock = mock_next_invocation_request(mock_response) + assert_next_invocation(get_mock, '') + assert_mock get_mock + end + + def test_next_invocation_with_null_tenant_id_header + mock_response = mock_next_invocation_response() + mock_response['Lambda-Runtime-Aws-Tenant-Id'] = nil + + get_mock = mock_next_invocation_request(mock_response) + assert_next_invocation(get_mock, nil) + assert_mock get_mock + end end diff --git a/test/unit/resources/runtime_handlers/ctx.rb b/test/unit/resources/runtime_handlers/ctx.rb index d577572..6517a8d 100644 --- a/test/unit/resources/runtime_handlers/ctx.rb +++ b/test/unit/resources/runtime_handlers/ctx.rb @@ -9,7 +9,8 @@ def get_context(event:,context:) log_group_name: context.log_group_name, log_stream_name: context.log_stream_name, memory_limit_in_mb: context.memory_limit_in_mb, - function_version: context.function_version + function_version: context.function_version, + tenant_id: context.tenant_id } end