diff --git a/apm-lambda-extension/main_test.go b/apm-lambda-extension/main_test.go index 8355f9f3..28e66d26 100644 --- a/apm-lambda-extension/main_test.go +++ b/apm-lambda-extension/main_test.go @@ -20,12 +20,8 @@ package main import ( "bytes" "context" - e2eTesting "elastic/apm-lambda-extension/e2e-testing" - "elastic/apm-lambda-extension/extension" - "elastic/apm-lambda-extension/logsapi" "encoding/json" "fmt" - "github.com/stretchr/testify/suite" "io/ioutil" "net" "net/http" @@ -36,6 +32,10 @@ import ( "testing" "time" + e2eTesting "elastic/apm-lambda-extension/e2e-testing" + "elastic/apm-lambda-extension/extension" + "elastic/apm-lambda-extension/logsapi" + "github.com/google/uuid" "github.com/stretchr/testify/assert" ) @@ -83,9 +83,7 @@ type ApmInfo struct { Version string `json:"version"` } -func initMockServers(eventsChannel chan MockEvent) (*httptest.Server, *httptest.Server, *MockServerInternals, *MockServerInternals) { - - // Mock APM Server +func newMockApmServer(t *testing.T) (*MockServerInternals, *httptest.Server) { var apmServerInternals MockServerInternals apmServerInternals.WaitForUnlockSignal = true apmServerInternals.UnlockSignalChannel = make(chan struct{}) @@ -137,16 +135,15 @@ func initMockServers(eventsChannel chan MockEvent) (*httptest.Server, *httptest. } } })) - if err := os.Setenv("ELASTIC_APM_LAMBDA_APM_SERVER", apmServer.URL); err != nil { - extension.Log.Fatalf("Could not set environment variable : %v", err) - return nil, nil, nil, nil - } - if err := os.Setenv("ELASTIC_APM_SECRET_TOKEN", "none"); err != nil { - extension.Log.Fatalf("Could not set environment variable : %v", err) - return nil, nil, nil, nil - } - // Mock Lambda Server + t.Setenv("ELASTIC_APM_LAMBDA_APM_SERVER", apmServer.URL) + t.Setenv("ELASTIC_APM_SECRET_TOKEN", "none") + + t.Cleanup(func() { apmServer.Close() }) + return &apmServerInternals, apmServer +} + +func newMockLambdaServer(t *testing.T, eventsChannel chan MockEvent) *MockServerInternals { var lambdaServerInternals MockServerInternals lambdaServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { switch r.RequestURI { @@ -185,10 +182,7 @@ func initMockServers(eventsChannel chan MockEvent) (*httptest.Server, *httptest. slicedLambdaURL := strings.Split(lambdaServer.URL, "//") strippedLambdaURL := slicedLambdaURL[1] - if err := os.Setenv("AWS_LAMBDA_RUNTIME_API", strippedLambdaURL); err != nil { - extension.Log.Fatalf("Could not set environment variable : %v", err) - return nil, nil, nil, nil - } + t.Setenv("AWS_LAMBDA_RUNTIME_API", strippedLambdaURL) extensionClient = extension.NewClient(os.Getenv("AWS_LAMBDA_RUNTIME_API")) // Find unused port for the extension to listen to @@ -197,12 +191,23 @@ func initMockServers(eventsChannel chan MockEvent) (*httptest.Server, *httptest. extension.Log.Errorf("Could not find free port for the extension to listen on : %v", err) extensionPort = 8200 } - if err = os.Setenv("ELASTIC_APM_DATA_RECEIVER_SERVER_PORT", fmt.Sprint(extensionPort)); err != nil { - extension.Log.Fatalf("Could not set environment variable : %v", err) - return nil, nil, nil, nil - } + t.Setenv("ELASTIC_APM_DATA_RECEIVER_SERVER_PORT", fmt.Sprint(extensionPort)) + + t.Cleanup(func() { lambdaServer.Close() }) + return &lambdaServerInternals +} + +// TODO : Move logger out of extension package and stop using it as a package-level variable +func initLogLevel(t *testing.T, logLevel string) { + t.Setenv("ELASTIC_APM_LOG_LEVEL", logLevel) +} - return lambdaServer, apmServer, &apmServerInternals, &lambdaServerInternals +func newTestStructs(t *testing.T) chan MockEvent { + http.DefaultServeMux = new(http.ServeMux) + _, cancel := context.WithCancel(context.Background()) + t.Cleanup(func() { cancel() }) + eventsChannel := make(chan MockEvent, 100) + return eventsChannel } func processMockEvent(currId string, event MockEvent, extensionPort string, internals *MockServerInternals) { @@ -335,199 +340,222 @@ func eventQueueGenerator(inputQueue []MockEvent, eventsChannel chan MockEvent) { } } -// TESTS -type MainUnitTestsSuite struct { - suite.Suite - eventsChannel chan MockEvent - lambdaServer *httptest.Server - apmServer *httptest.Server - apmServerInternals *MockServerInternals - lambdaServerInternals *MockServerInternals - ctx context.Context - cancel context.CancelFunc -} - -func TestMainUnitTestsSuite(t *testing.T) { - suite.Run(t, new(MainUnitTestsSuite)) -} - -// This function executes before each test case -func (suite *MainUnitTestsSuite) SetupTest() { - if err := os.Setenv("ELASTIC_APM_LOG_LEVEL", "trace"); err != nil { - suite.T().Fail() - } - suite.ctx, suite.cancel = context.WithCancel(context.Background()) - http.DefaultServeMux = new(http.ServeMux) - suite.eventsChannel = make(chan MockEvent, 100) - suite.lambdaServer, suite.apmServer, suite.apmServerInternals, suite.lambdaServerInternals = initMockServers(suite.eventsChannel) -} - -// This function executes after each test case -func (suite *MainUnitTestsSuite) TearDownTest() { - suite.lambdaServer.Close() - suite.apmServer.Close() - suite.cancel() -} - // TestStandardEventsChain checks a nominal sequence of events (fast APM server, only one standard event) -func (suite *MainUnitTestsSuite) TestStandardEventsChain() { +func TestStandardEventsChain(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeStandard, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.True(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(TimelyResponse))) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.Contains(t, apmServerInternals.Data, TimelyResponse) } // TestFlush checks if the flushed param does not cause a panic or an unexpected behavior -func (suite *MainUnitTestsSuite) TestFlush() { +func TestFlush(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeStandardFlush, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.True(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(TimelyResponse))) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.Contains(t, apmServerInternals.Data, TimelyResponse) } // TestLateFlush checks if there is no race condition between RuntimeDone and AgentDone // The test is built so that the AgentDone signal is received after RuntimeDone, which causes the next event to be interrupted. -func (suite *MainUnitTestsSuite) TestLateFlush() { +func TestLateFlush(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeLateFlush, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 5}, {Type: InvokeStandard, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.True(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(TimelyResponse+TimelyResponse))) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.Contains(t, apmServerInternals.Data, TimelyResponse+TimelyResponse) } // TestWaitGroup checks if there is no race condition between the main waitgroups (issue #128) -func (suite *MainUnitTestsSuite) TestWaitGroup() { +func TestWaitGroup(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeWaitgroupsRace, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 500}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.True(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(TimelyResponse))) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.Contains(t, apmServerInternals.Data, TimelyResponse) } // TestAPMServerDown tests that main does not panic nor runs indefinitely when the APM server is inactive. -func (suite *MainUnitTestsSuite) TestAPMServerDown() { - suite.apmServer.Close() +func TestAPMServerDown(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, apmServer := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + + apmServer.Close() eventsChain := []MockEvent{ {Type: InvokeStandard, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.False(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(TimelyResponse))) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.NotContains(t, apmServerInternals.Data, TimelyResponse) } // TestAPMServerHangs tests that main does not panic nor runs indefinitely when the APM server does not respond. -func (suite *MainUnitTestsSuite) TestAPMServerHangs() { +func TestAPMServerHangs(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeStandard, APMServerBehavior: Hangs, ExecutionDuration: 1, Timeout: 500}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.False(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(Hangs))) - suite.apmServerInternals.UnlockSignalChannel <- struct{}{} + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.NotContains(t, apmServerInternals.Data, Hangs) + apmServerInternals.UnlockSignalChannel <- struct{}{} } // TestAPMServerRecovery tests a scenario where the APM server recovers after hanging. // The default forwarder timeout is 3 seconds, so we wait 5 seconds until we unlock that hanging state. // Hence, the APM server is waked up just in time to process the TimelyResponse data frame. -func (suite *MainUnitTestsSuite) TestAPMServerRecovery() { - if err := os.Setenv("ELASTIC_APM_DATA_FORWARDER_TIMEOUT_SECONDS", "1"); err != nil { - suite.T().Fail() - } +func TestAPMServerRecovery(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + + t.Setenv("ELASTIC_APM_DATA_FORWARDER_TIMEOUT_SECONDS", "1") eventsChain := []MockEvent{ {Type: InvokeStandard, APMServerBehavior: Hangs, ExecutionDuration: 1, Timeout: 5}, {Type: InvokeStandard, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) + eventQueueGenerator(eventsChain, eventsChannel) go func() { time.Sleep(2500 * time.Millisecond) // Cannot multiply time.Second by a float - suite.apmServerInternals.UnlockSignalChannel <- struct{}{} + apmServerInternals.UnlockSignalChannel <- struct{}{} }() - assert.NotPanics(suite.T(), main) - assert.True(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(Hangs))) - assert.True(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(TimelyResponse))) - if err := os.Setenv("ELASTIC_APM_DATA_FORWARDER_TIMEOUT_SECONDS", ""); err != nil { - suite.T().Fail() - } + assert.NotPanics(t, main) + assert.Contains(t, apmServerInternals.Data, Hangs) + assert.Contains(t, apmServerInternals.Data, TimelyResponse) + t.Setenv("ELASTIC_APM_DATA_FORWARDER_TIMEOUT_SECONDS", "") } // TestGracePeriodHangs verifies that the WaitforGracePeriod goroutine ends when main() ends. // This can be checked by asserting that apmTransportStatus is pending after the execution of main. -func (suite *MainUnitTestsSuite) TestGracePeriodHangs() { +func TestGracePeriodHangs(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeStandard, APMServerBehavior: Hangs, ExecutionDuration: 1, Timeout: 500}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) time.Sleep(100 * time.Millisecond) - suite.apmServerInternals.UnlockSignalChannel <- struct{}{} + apmServerInternals.UnlockSignalChannel <- struct{}{} } // TestAPMServerCrashesDuringExecution tests that main does not panic nor runs indefinitely when the APM server crashes // during execution. -func (suite *MainUnitTestsSuite) TestAPMServerCrashesDuringExecution() { +func TestAPMServerCrashesDuringExecution(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeStandard, APMServerBehavior: Crashes, ExecutionDuration: 1, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.False(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(Crashes))) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.NotContains(t, apmServerInternals.Data, Crashes) } // TestFullChannel checks that an overload of APM data chunks is handled correctly, events dropped beyond the 100th one // if no space left in channel, no panic, no infinite hanging. -func (suite *MainUnitTestsSuite) TestFullChannel() { +func TestFullChannel(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeMultipleTransactionsOverload, APMServerBehavior: TimelyResponse, ExecutionDuration: 0.1, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.True(suite.T(), strings.Contains(suite.apmServerInternals.Data, string(TimelyResponse))) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.Contains(t, apmServerInternals.Data, TimelyResponse) } // TestFullChannelSlowAPMServer tests what happens when the APM Data channel is full and the APM server is slow // (send strategy : background) -func (suite *MainUnitTestsSuite) TestFullChannelSlowAPMServer() { - if err := os.Setenv("ELASTIC_APM_SEND_STRATEGY", "background"); err != nil { - suite.T().Fail() - } +func TestFullChannelSlowAPMServer(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + newMockApmServer(t) + newMockLambdaServer(t, eventsChannel) + + t.Setenv("ELASTIC_APM_SEND_STRATEGY", "background") eventsChain := []MockEvent{ {Type: InvokeMultipleTransactionsOverload, APMServerBehavior: SlowResponse, ExecutionDuration: 0.01, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) // The test should not hang - if err := os.Setenv("ELASTIC_APM_SEND_STRATEGY", "syncflush"); err != nil { - suite.T().Fail() - } + t.Setenv("ELASTIC_APM_SEND_STRATEGY", "syncflush") } // TestInfoRequest checks if the extension is able to retrieve APM server info (/ endpoint) (fast APM server, only one standard event) -func (suite *MainUnitTestsSuite) TestInfoRequest() { +func TestInfoRequest(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + newMockApmServer(t) + lambdaServerInternals := newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeStandardInfo, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 5}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.True(suite.T(), strings.Contains(suite.lambdaServerInternals.Data, "7814d524d3602e70b703539c57568cba6964fc20")) + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.Contains(t, lambdaServerInternals.Data, "7814d524d3602e70b703539c57568cba6964fc20") } // TestInfoRequest checks if the extension times out when unable to retrieve APM server info (/ endpoint) -func (suite *MainUnitTestsSuite) TestInfoRequestHangs() { +func TestInfoRequestHangs(t *testing.T) { + initLogLevel(t, "trace") + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t) + lambdaServerInternals := newMockLambdaServer(t, eventsChannel) + eventsChain := []MockEvent{ {Type: InvokeStandardInfo, APMServerBehavior: Hangs, ExecutionDuration: 1, Timeout: 500}, } - eventQueueGenerator(eventsChain, suite.eventsChannel) - assert.NotPanics(suite.T(), main) - assert.False(suite.T(), strings.Contains(suite.lambdaServerInternals.Data, "7814d524d3602e70b703539c57568cba6964fc20")) - suite.apmServerInternals.UnlockSignalChannel <- struct{}{} + eventQueueGenerator(eventsChain, eventsChannel) + assert.NotPanics(t, main) + assert.NotContains(t, lambdaServerInternals.Data, "7814d524d3602e70b703539c57568cba6964fc20") + apmServerInternals.UnlockSignalChannel <- struct{}{} }