Skip to content

Commit b003ddb

Browse files
geroplgeropl
committed
Implement lambda-handler tests
Tool: gitpod/catfood.gitpod.cloud
1 parent c8a3521 commit b003ddb

File tree

4 files changed

+214
-82
lines changed

4 files changed

+214
-82
lines changed

‎gitpod-network-check/cmd/lambda_handler.go‎

Lines changed: 64 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -1,113 +1,95 @@
11
package cmd
22

33
import (
4-
"encoding/json"
4+
"context"
55
"fmt"
6-
"io"
76
"net/http"
87
"os"
98
"time"
109

10+
"github.com/aws/aws-lambda-go/lambda"
1111
log "github.com/sirupsen/logrus"
1212
"github.com/spf13/cobra"
1313

1414
"github.com/gitpod-io/enterprise-deployment-toolkit/gitpod-network-check/pkg/lambda_types"
1515
)
1616

17-
var lambdaHandlerCmd = &cobra.Command{
18-
Use: "lambda-handler",
19-
Short: "Internal command to execute network checks within AWS Lambda (reads JSON request from stdin, writes JSON response to stdout)",
20-
Hidden: true, // Hide this command from user help output
21-
PersistentPreRun: func(cmd *cobra.Command, args []string) {
22-
// override parent, as we don't care about the config or other flags
23-
},
24-
RunE: func(cmd *cobra.Command, args []string) error {
25-
// Lambda environment might not have sophisticated logging setup, print directly
26-
fmt.Fprintln(os.Stderr, "Lambda Handler: Starting execution.")
27-
28-
// Read request payload from stdin
29-
stdinBytes, err := io.ReadAll(os.Stdin)
30-
if err != nil {
31-
fmt.Fprintf(os.Stderr, "Lambda Handler: Error reading stdin: %v\n", err)
32-
return fmt.Errorf("error reading stdin: %w", err)
33-
}
17+
// Core logic for handling the Lambda event
18+
// This function is called by the aws-lambda-go library.
19+
func handleLambdaEvent(ctx context.Context, request lambda_types.CheckRequest) (lambda_types.CheckResponse, error) {
20+
log.Infof("Lambda Handler: Received check request for %d endpoints.", len(request.Endpoints))
3421

35-
var request lambda_types.CheckRequest
36-
err = json.Unmarshal(stdinBytes, &request)
37-
if err != nil {
38-
fmt.Fprintf(os.Stderr, "Lambda Handler: Error unmarshalling request JSON: %v\n", err)
39-
fmt.Fprintf(os.Stderr, "Lambda Handler: Received input: %s\n", string(stdinBytes))
40-
return fmt.Errorf("error unmarshalling request: %w", err)
41-
}
22+
response := lambda_types.CheckResponse{
23+
Results: make(map[string]lambda_types.CheckResult),
24+
}
4225

43-
fmt.Fprintf(os.Stderr, "Lambda Handler: Received check request for %d endpoints.\n", len(request.Endpoints))
26+
client := &http.Client{
27+
Timeout: 10 * time.Second, // Consider making this configurable if needed
28+
}
4429

45-
response := lambda_types.CheckResponse{
46-
Results: make(map[string]lambda_types.CheckResult),
47-
}
30+
// Perform checks
31+
for name, url := range request.Endpoints {
32+
log.Debugf("Lambda Handler: Checking endpoint: %s (%s)", name, url)
4833

49-
client := &http.Client{
50-
Timeout: 10 * time.Second, // Slightly longer timeout for Lambda environment?
34+
// Use the context provided by the Lambda runtime
35+
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
36+
if err != nil {
37+
response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("failed to create request: %v", err)}
38+
log.Warnf(" -> Failed (request creation): %v", err)
39+
continue
5140
}
5241

53-
// Perform checks (similar logic to the previous dedicated handler)
54-
for name, url := range request.Endpoints {
55-
fmt.Fprintf(os.Stderr, "Lambda Handler: Checking endpoint: %s (%s)\n", name, url)
56-
// Use context from command if needed, otherwise background context is fine here
57-
req, err := http.NewRequestWithContext(cmd.Context(), "GET", url, nil)
58-
if err != nil {
59-
response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("failed to create request: %v", err)}
60-
fmt.Fprintf(os.Stderr, " -> Failed (request creation): %v\n", err)
61-
continue
62-
}
63-
64-
resp, err := client.Do(req)
65-
if err != nil {
66-
response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("HTTP request failed: %v", err)}
67-
fmt.Fprintf(os.Stderr, " -> Failed (HTTP request): %v\n", err)
42+
resp, err := client.Do(req)
43+
if err != nil {
44+
response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("HTTP request failed: %v", err)}
45+
log.Warnf(" -> Failed (HTTP request): %v", err)
46+
} else {
47+
resp.Body.Close() // Ensure body is closed
48+
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
49+
response.Results[name] = lambda_types.CheckResult{Success: true}
50+
log.Debugf(" -> Success (Status: %d)", resp.StatusCode)
6851
} else {
69-
resp.Body.Close() // Ensure body is closed
70-
if resp.StatusCode >= 200 && resp.StatusCode < 300 {
71-
response.Results[name] = lambda_types.CheckResult{Success: true}
72-
fmt.Fprintf(os.Stderr, " -> Success (Status: %d)\n", resp.StatusCode)
73-
} else {
74-
response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("unexpected status code: %d", resp.StatusCode)}
75-
fmt.Fprintf(os.Stderr, " -> Failed (Status: %d)\n", resp.StatusCode)
76-
}
52+
response.Results[name] = lambda_types.CheckResult{Success: false, Error: fmt.Sprintf("unexpected status code: %d", resp.StatusCode)}
53+
log.Warnf(" -> Failed (Status: %d)", resp.StatusCode)
7754
}
7855
}
56+
}
7957

80-
// Marshal response payload to stdout
81-
responseBytes, err := json.Marshal(response)
82-
if err != nil {
83-
fmt.Fprintf(os.Stderr, "Lambda Handler: Error marshalling response JSON: %v\n", err)
84-
return fmt.Errorf("error marshalling response: %w", err)
85-
}
86-
87-
_, err = fmt.Fprint(os.Stdout, string(responseBytes))
88-
if err != nil {
89-
fmt.Fprintf(os.Stderr, "Lambda Handler: Error writing response to stdout: %v\n", err)
90-
return fmt.Errorf("error writing response: %w", err)
91-
}
58+
log.Info("Lambda Handler: Check processing complete.")
59+
// The lambda library handles marshalling the response and deals with errors.
60+
// We return the response struct and nil error if processing logic itself didn't fail critically.
61+
return response, nil
62+
}
9263

93-
fmt.Fprintln(os.Stderr, "Lambda Handler: Execution complete.")
94-
return nil
64+
// lambdaHandlerCmd is the Cobra command invoked when the binary is run with the "lambda-handler" argument.
65+
// This happens inside the AWS Lambda environment via the bootstrap script.
66+
var lambdaHandlerCmd = &cobra.Command{
67+
Use: "lambda-handler",
68+
Short: "Internal command used by AWS Lambda runtime to execute network checks",
69+
Hidden: true, // Hide this command from user help output
70+
PersistentPreRun: func(cmd *cobra.Command, args []string) {
71+
// override parent, as we don't care about the config or other flags when run by lambda
72+
// Ensure logs go to stderr (Lambda standard)
73+
log.SetOutput(os.Stderr)
74+
// Optionally set log level from env var if needed, e.g., os.Getenv("LOG_LEVEL")
75+
// Consider setting a default level appropriate for Lambda execution.
76+
log.SetLevel(log.InfoLevel) // Example: Set a default level
77+
},
78+
RunE: func(cmd *cobra.Command, args []string) error {
79+
// The aws-lambda-go library takes over execution when lambda.Start is called.
80+
// It handles reading events, invoking the handler, and writing responses.
81+
log.Info("Lambda Handler: Starting AWS Lambda handler loop.")
82+
lambda.Start(handleLambdaEvent)
83+
// lambda.Start blocks and never returns unless there's a critical error during initialization
84+
log.Error("Lambda Handler: lambda.Start returned unexpectedly (should not happen)")
85+
return fmt.Errorf("lambda.Start returned unexpectedly")
9586
},
96-
// Disable flag parsing for this internal command as it gets input via stdin
87+
// Disable flag parsing for this internal command as input comes from Lambda event payload
9788
DisableFlagParsing: true,
9889
}
9990

10091
func init() {
101-
// Note: We don't add this to networkCheckCmd directly in init() here
102-
// because it might interfere with normal flag parsing if not careful.
103-
// It will be added in the main Execute() function or similar central place.
104-
// For now, just define the command struct.
105-
// We also need to ensure logging doesn't interfere with stdout JSON output.
106-
// Maybe configure logging to stderr specifically for this command?
107-
lambdaHandlerCmd.PersistentPreRun = func(cmd *cobra.Command, args []string) {
108-
// Ensure logs go to stderr for this command to keep stdout clean for JSON
109-
log.SetOutput(os.Stderr)
110-
}
111-
112-
NetworkCheckCmd.AddCommand(lambdaHandlerCmd) // Register the hidden lambda handler command
92+
// Register the hidden lambda handler command
93+
// It's invoked by the Lambda runtime via the bootstrap script
94+
NetworkCheckCmd.AddCommand(lambdaHandlerCmd)
11395
}

‎gitpod-network-check/cmd/lambda_handler_test.go‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
package cmd
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"net/http"
7+
"net/http/httptest"
8+
"testing"
9+
10+
"github.com/gitpod-io/enterprise-deployment-toolkit/gitpod-network-check/pkg/lambda_types"
11+
cmp "github.com/google/go-cmp/cmp" // Added for deep comparison
12+
log "github.com/sirupsen/logrus"
13+
)
14+
15+
// TestHandleLambdaEvent tests the core logic within the Lambda handler function.
16+
func TestHandleLambdaEvent(t *testing.T) {
17+
18+
tests := []struct {
19+
name string
20+
request lambda_types.CheckRequest
21+
expectedResp lambda_types.CheckResponse
22+
}{
23+
{
24+
name: "successful http check",
25+
request: lambda_types.CheckRequest{
26+
Endpoints: map[string]string{
27+
"example_http": "http://example.com", // Use http to avoid cert issues in test
28+
},
29+
},
30+
expectedResp: lambda_types.CheckResponse{
31+
Results: map[string]lambda_types.CheckResult{
32+
"example_http": {Success: true},
33+
},
34+
},
35+
},
36+
{
37+
name: "successful https check",
38+
request: lambda_types.CheckRequest{
39+
Endpoints: map[string]string{
40+
"example_https": "https://example.com",
41+
},
42+
},
43+
expectedResp: lambda_types.CheckResponse{
44+
Results: map[string]lambda_types.CheckResult{
45+
"example_https": {Success: true},
46+
},
47+
},
48+
},
49+
{
50+
name: "failed http check - 404",
51+
// Assuming httpbin gives a 404 for this path
52+
request: lambda_types.CheckRequest{
53+
Endpoints: map[string]string{
54+
"httpbin_404": "http://httpbin.org/status/404",
55+
},
56+
},
57+
expectedResp: lambda_types.CheckResponse{
58+
Results: map[string]lambda_types.CheckResult{
59+
"httpbin_404": {Success: false, Error: "unexpected status code: 404"},
60+
},
61+
},
62+
},
63+
{
64+
name: "failed http check - connection refused",
65+
// Use a port likely not open on localhost
66+
request: lambda_types.CheckRequest{
67+
Endpoints: map[string]string{
68+
"localhost_conn_refused": "http://127.0.0.1:1",
69+
},
70+
},
71+
expectedResp: lambda_types.CheckResponse{
72+
Results: map[string]lambda_types.CheckResult{
73+
// Error message might vary slightly depending on OS/network stack - adjust if needed after running
74+
"localhost_conn_refused": {Success: false, Error: "HTTP request failed: Get \"http://127.0.0.1:1\": dial tcp 127.0.0.1:1: connect: connection refused"},
75+
},
76+
},
77+
},
78+
{
79+
name: "multiple endpoints - mix success and failure",
80+
request: lambda_types.CheckRequest{
81+
Endpoints: map[string]string{
82+
"example_ok": "https://example.com",
83+
"httpbin_404": "http://httpbin.org/status/404",
84+
"conn_refused": "http://127.0.0.1:1",
85+
},
86+
},
87+
expectedResp: lambda_types.CheckResponse{
88+
Results: map[string]lambda_types.CheckResult{
89+
"example_ok": {Success: true},
90+
"httpbin_404": {Success: false, Error: "unexpected status code: 404"},
91+
// Error message might vary slightly depending on OS/network stack - adjust if needed after running
92+
"conn_refused": {Success: false, Error: "HTTP request failed: Get \"http://127.0.0.1:1\": dial tcp 127.0.0.1:1: connect: connection refused"},
93+
},
94+
},
95+
},
96+
}
97+
98+
// Setup mock HTTP server for reliable testing of external URLs
99+
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
100+
// Check the path to determine the response, as Host will be the mock server's address
101+
if r.URL.Path == "/status/404" {
102+
w.WriteHeader(http.StatusNotFound)
103+
fmt.Fprintln(w, "Not Found")
104+
} else if r.URL.Path == "/" { // Assume requests to the root are for the "success" case
105+
w.WriteHeader(http.StatusOK)
106+
fmt.Fprintln(w, "OK")
107+
} else {
108+
// Log unexpected paths to help debug test failures
109+
log.Warnf("Mock server received request for unexpected path: %s", r.URL.Path)
110+
w.WriteHeader(http.StatusInternalServerError)
111+
fmt.Fprintln(w, "Mock server default - unexpected path")
112+
}
113+
}))
114+
defer mockServer.Close()
115+
116+
// Replace external URLs in test cases with mock server URL
117+
// This makes tests faster and more reliable (no external network dependency)
118+
for i := range tests {
119+
newEndpoints := make(map[string]string)
120+
for name, url := range tests[i].request.Endpoints {
121+
switch url {
122+
case "http://example.com", "https://example.com":
123+
newEndpoints[name] = mockServer.URL
124+
case "http://httpbin.org/status/404":
125+
newEndpoints[name] = mockServer.URL + "/status/404"
126+
default:
127+
newEndpoints[name] = url // Keep internal/localhost URLs as is
128+
}
129+
}
130+
tests[i].request.Endpoints = newEndpoints
131+
}
132+
133+
for _, tt := range tests {
134+
tt := tt // Capture range variable
135+
t.Run(tt.name, func(t *testing.T) {
136+
actualResp, err := handleLambdaEvent(context.Background(), tt.request)
137+
if err != nil {
138+
t.Errorf("handleLambdaEvent returned an unexpected error: %v", err)
139+
}
140+
141+
if diff := cmp.Diff(tt.expectedResp, actualResp); diff != "" {
142+
t.Errorf("handleLambdaEvent response mismatch (-want +got):\n%s", diff)
143+
}
144+
})
145+
}
146+
}

‎gitpod-network-check/go.mod‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ module github.com/gitpod-io/enterprise-deployment-toolkit/gitpod-network-check
33
go 1.23.1
44

55
require (
6+
github.com/aws/aws-lambda-go v1.47.0
67
github.com/aws/aws-sdk-go-v2 v1.36.3
78
github.com/aws/aws-sdk-go-v2/config v1.27.6
89
github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs v1.47.1
@@ -11,6 +12,7 @@ require (
1112
github.com/aws/aws-sdk-go-v2/service/lambda v1.71.0
1213
github.com/aws/aws-sdk-go-v2/service/ssm v1.49.1
1314
github.com/aws/smithy-go v1.22.2
15+
github.com/google/go-cmp v0.6.0
1416
github.com/sirupsen/logrus v1.9.3
1517
github.com/spf13/cobra v1.8.0
1618
github.com/spf13/pflag v1.0.5

‎gitpod-network-check/go.sum‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
github.com/aws/aws-lambda-go v1.47.0 h1:0H8s0vumYx/YKs4sE7YM0ktwL2eWse+kfopsRI1sXVI=
2+
github.com/aws/aws-lambda-go v1.47.0/go.mod h1:dpMpZgvWx5vuQJfBt0zqBha60q7Dd7RfgJv23DymV8A=
13
github.com/aws/aws-sdk-go-v2 v1.36.3 h1:mJoei2CxPutQVxaATCzDUjcZEjVRdpsiiXi2o38yqWM=
24
github.com/aws/aws-sdk-go-v2 v1.36.3/go.mod h1:LLXuLpgzEbD766Z5ECcRmi8AzSwfZItDtmABVkRLGzg=
35
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.10 h1:zAybnyUQXIZ5mok5Jqwlf58/TFE7uvd3IAsa1aF9cXs=

0 commit comments

Comments
 (0)