feat: add tenant-id to lambda context (#43)

Co-authored-by: Ramisa Alam <[email protected]>

Author: ramisa2108

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
 

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

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