diff --git a/ridge.go b/ridge.go index b96abd5..adba74d 100644 --- a/ridge.go +++ b/ridge.go @@ -181,22 +181,33 @@ func (r *Ridge) Run() { // RunWithContext runs http handler on AWS Lambda runtime or net/http's server with context. func (r *Ridge) RunWithContext(ctx context.Context) { - if IsOnLambdaRuntime() { - r.runOnLambdaRuntime(ctx) + if AsLambdaHandler() { + r.runAsLambdaHandler(ctx) } else { + // If it is not running on the AWS Lambda runtime or running as a Lambda extension, + // runs a net/http server. r.runOnNetHTTPServer(ctx) } } -// IsOnLambdaRuntime returns true if running on AWS Lambda runtime (excludes extensions). +// OnLambdaRuntime returns true if running on AWS Lambda runtime // - AWS_EXECUTION_ENV is set on AWS Lambda runtime (go1.x) // - AWS_LAMBDA_RUNTIME_API is set on custom runtime (provided.*) -// - _HANDLER is not set on AWS Lambda extension -func IsOnLambdaRuntime() bool { - return (strings.HasPrefix(os.Getenv("AWS_EXECUTION_ENV"), "AWS_Lambda") || os.Getenv("AWS_LAMBDA_RUNTIME_API") != "") && os.Getenv("_HANDLER") != "" +func OnLambdaRuntime() bool { + return (strings.HasPrefix(os.Getenv("AWS_EXECUTION_ENV"), "AWS_Lambda") || os.Getenv("AWS_LAMBDA_RUNTIME_API") != "") } -func (r *Ridge) runOnLambdaRuntime(ctx context.Context) { +// AsLambdaExtension returns true if running on AWS Lambda runtime and run as a Lambda extension +func AsLambdaExtension() bool { + return OnLambdaRuntime() && os.Getenv("_HANDLER") == "" +} + +// AsLambdaHandler returns true if running on AWS Lambda runtime and run as a Lambda handler +func AsLambdaHandler() bool { + return OnLambdaRuntime() && os.Getenv("_HANDLER") != "" +} + +func (r *Ridge) runAsLambdaHandler(ctx context.Context) { handler := func(event json.RawMessage) (interface{}, error) { req, err := r.RequestBuilder(event) if err != nil { diff --git a/ridge_test.go b/ridge_test.go index 61d3314..cdc3306 100644 --- a/ridge_test.go +++ b/ridge_test.go @@ -6,19 +6,21 @@ import ( "github.com/fujiwara/ridge" ) -func TestIsOnLambdaRuntime(t *testing.T) { +func TestRuntimeEnvironments(t *testing.T) { tests := []struct { name string awsExecutionEnv string awsLambdaRuntimeAPI string handler string - expected bool + onLambdaRuntime bool + asLambdaHandler bool + asLambdaExtension bool }{ - {"On AWS Lambda go runtime", "AWS_Lambda_go1.x", "", "handler", true}, - {"On AWS Lambda custom runtime", "", "http://localhost:8080", "handler", true}, - {"On AWS Lambda extension with go runtime", "AWS_Lambda_go1.x", "", "", false}, - {"On AWS Lambda extension with custom runtime", "", "http://localhost:8080", "", false}, - {"Not on AWS Lambda", "", "", "", false}, + {"On AWS Lambda go runtime handler", "AWS_Lambda_go1.x", "", "handler", true, true, false}, + {"On AWS Lambda custom runtime handler", "", "http://localhost:8080", "handler", true, true, false}, + {"On AWS Lambda extension with go runtime", "AWS_Lambda_go1.x", "", "", true, false, true}, + {"On AWS Lambda extension with custom runtime", "", "http://localhost:8080", "", true, false, true}, + {"Not on AWS Lambda", "", "", "", false, false, false}, } for _, tt := range tests { @@ -27,8 +29,14 @@ func TestIsOnLambdaRuntime(t *testing.T) { t.Setenv("AWS_LAMBDA_RUNTIME_API", tt.awsLambdaRuntimeAPI) t.Setenv("_HANDLER", tt.handler) - if got := ridge.IsOnLambdaRuntime(); got != tt.expected { - t.Errorf("IsOnLambdaRuntime() = %v; want %v", got, tt.expected) + if ridge.OnLambdaRuntime() != tt.onLambdaRuntime { + t.Errorf("OnLambdaRuntime() = %v, want %v", ridge.OnLambdaRuntime(), tt.onLambdaRuntime) + } + if ridge.AsLambdaHandler() != tt.asLambdaHandler { + t.Errorf("AsLambdaHandler() = %v, want %v", ridge.AsLambdaHandler(), tt.asLambdaHandler) + } + if ridge.AsLambdaExtension() != tt.asLambdaExtension { + t.Errorf("AsLambdaExtension() = %v, want %v", ridge.AsLambdaExtension(), tt.asLambdaExtension) } }) }