[collector] add tests for lifecycle manager (aws-observability#410)

* [collector] add tests for lifecycle manager

Improved the error handling in processEvents along the way.

Signed-off-by: Alex Boten <[email protected]>

* move collectorWrapper interface

Signed-off-by: Alex Boten <[email protected]>

Signed-off-by: Alex Boten <[email protected]>

Author: Alex Boten

collector/go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/golang-collections/go-datastructures v0.0.0-20150211160725-59788d5eb259
 	github.com/open-telemetry/opentelemetry-collector-contrib/confmap/provider/s3provider v0.68.0
 	github.com/open-telemetry/opentelemetry-lambda/collector/lambdacomponents v0.0.0
+	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.1
 	go.opentelemetry.io/collector v0.68.0
 	go.opentelemetry.io/collector/component v0.68.0
@@ -41,6 +42,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.13.11 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.17.7 // indirect
 	github.com/aws/smithy-go v1.13.5 // indirect
+	github.com/benbjohnson/clock v1.3.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cenkalti/backoff/v4 v4.2.0 // indirect
 	github.com/cespare/xxhash/v2 v2.1.2 // indirect

collector/go.sum

@@ -409,6 +409,7 @@ github.com/aws/smithy-go v1.9.0/go.mod h1:SObp3lf9smib00L/v3U2eAKG8FyQ7iLrJnQiAm
 github.com/aws/smithy-go v1.13.5 h1:hgz0X/DX0dGqTYpGALqXJoRKRj5oQ7150i5FdTePzO8=
 github.com/aws/smithy-go v1.13.5/go.mod h1:Tg+OJXh4MB2R/uN61Ko2f6hTZwB/ZYGOtib8J3gBHzA=
 github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
+github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/benbjohnson/immutable v0.2.1/go.mod h1:uc6OHo6PN2++n98KHLxW8ef4W42ylHiQSENghE1ezxI=
 github.com/benbjohnson/tmpl v1.0.0/go.mod h1:igT620JFIi44B6awvU9IsDhR77IXWtFigTLil/RPdps=
 github.com/beorn7/perks v0.0.0-20160804104726-4c0e84591b9a/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=

collector/internal/collector/collector.go

@@ -33,8 +33,8 @@ import (
 	"go.uber.org/zap/zapcore"
 )
 
-// Collector implements the OtelcolRunner interfaces running a single otelcol as a go routine within the
-// same process as the test executor.
+// Collector runs a single otelcol as a go routine within the
+// same process as the executor.
 type Collector struct {
 	factories      component.Factories
 	configProvider otelcol.ConfigProvider

collector/internal/lifecycle/manager.go

@@ -27,16 +27,22 @@ import (
 	"github.com/open-telemetry/opentelemetry-lambda/collector/internal/extensionapi"
 	"github.com/open-telemetry/opentelemetry-lambda/collector/internal/telemetryapi"
 	"github.com/open-telemetry/opentelemetry-lambda/collector/lambdacomponents"
+	"github.com/pkg/errors"
 	"go.uber.org/zap"
 )
 
 var (
 	extensionName = filepath.Base(os.Args[0]) // extension name has to match the filename
 )
 
+type collectorWrapper interface {
+	Start(ctx context.Context) error
+	Stop() error
+}
+
 type manager struct {
 	logger          *zap.Logger
-	collector       *collector.Collector
+	collector       collectorWrapper
 	extensionClient *extensionapi.Client
 	listener        *telemetryapi.Listener
 	wg              sync.WaitGroup
@@ -77,7 +83,11 @@ func NewManager(ctx context.Context, logger *zap.Logger, version string) (contex
 		listener:        listener,
 	}
 
-	go lm.processEvents(ctx)
+	go func() {
+		if err := lm.processEvents(ctx); err != nil {
+			lm.logger.Warn("Failed to process events", zap.Error(err))
+		}
+	}()
 
 	factories, _ := lambdacomponents.Components()
 	lm.collector = collector.NewCollector(logger, factories, version)
@@ -96,21 +106,23 @@ func (lm *manager) Run(ctx context.Context) error {
 	return nil
 }
 
-func (lm *manager) processEvents(ctx context.Context) {
+func (lm *manager) processEvents(ctx context.Context) error {
 	lm.wg.Add(1)
 	defer lm.wg.Done()
 
 	for {
 		select {
 		case <-ctx.Done():
-			return
+			return nil
 		default:
 			lm.logger.Debug("Waiting for event...")
 			res, err := lm.extensionClient.NextEvent(ctx)
 			if err != nil {
 				lm.logger.Warn("error waiting for extension event", zap.Error(err))
-				lm.extensionClient.ExitError(ctx, fmt.Sprintf("error waiting for extension event: %v", err))
-				return
+				if _, exitErr := lm.extensionClient.ExitError(ctx, fmt.Sprintf("error waiting for extension event: %v", err)); exitErr != nil {
+					err = errors.Wrap(err, exitErr.Error())
+				}
+				return err
 			}
 
 			lm.logger.Debug("Received ", zap.Any("event :", res))
@@ -120,9 +132,11 @@ func (lm *manager) processEvents(ctx context.Context) {
 				lm.listener.Shutdown()
 				err = lm.collector.Stop()
 				if err != nil {
-					lm.extensionClient.ExitError(ctx, fmt.Sprintf("error stopping collector: %v", err))
+					if _, exitErr := lm.extensionClient.ExitError(ctx, fmt.Sprintf("error stopping collector: %v", err)); exitErr != nil {
+						err = errors.Wrap(err, exitErr.Error())
+					}
 				}
-				return
+				return err
 			}
 
 			err = lm.listener.Wait(ctx, res.RequestID)

collector/internal/lifecycle/manager_test.go

@@ -0,0 +1,142 @@
+// Copyright The OpenTelemetry Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package lifecycle
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+
+	"github.com/open-telemetry/opentelemetry-lambda/collector/internal/extensionapi"
+	"github.com/open-telemetry/opentelemetry-lambda/collector/internal/telemetryapi"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap/zaptest"
+)
+
+type MockCollector struct {
+	err error
+}
+
+func (c *MockCollector) Start(ctx context.Context) error {
+	return c.err
+}
+func (c *MockCollector) Stop() error {
+	return c.err
+}
+
+func TestRun(t *testing.T) {
+	logger := zaptest.NewLogger(t)
+	ctx := context.Background()
+	// test with an error
+	lm := manager{
+		collector:       &MockCollector{err: fmt.Errorf("test start error")},
+		logger:          logger,
+		extensionClient: extensionapi.NewClient(logger, ""),
+	}
+	require.Error(t, lm.Run(ctx))
+	// test with no waitgroup
+	lm = manager{
+		collector:       &MockCollector{},
+		logger:          logger,
+		extensionClient: extensionapi.NewClient(logger, ""),
+	}
+	require.NoError(t, lm.Run(ctx))
+	// test with waitgroup counter incremented
+	lm = manager{
+		collector:       &MockCollector{},
+		logger:          logger,
+		extensionClient: extensionapi.NewClient(logger, ""),
+	}
+	lm.wg.Add(1)
+	go func() {
+		require.NoError(t, lm.Run(ctx))
+	}()
+	lm.wg.Done()
+
+}
+
+func TestProcessEvents(t *testing.T) {
+	type test struct {
+		name           string
+		cancel         bool
+		err            error
+		serverResponse string
+		collectorError error
+	}
+	testCases := []test{
+		{
+			name:   "processEvents with context cancelled",
+			cancel: true,
+		},
+		{
+			name: "processEvents with error from extension API",
+			err:  fmt.Errorf("unexpected end of JSON input"),
+		},
+		{
+			name:           "processEvents with shutdown event received",
+			serverResponse: `{"time":"2006-01-02T15:04:05.000Z", "eventType":"SHUTDOWN", "record":{}}`,
+		},
+		{
+			name:           "processEvents with shutdown event received and collector error",
+			serverResponse: `{"time":"2006-01-02T15:04:05.000Z", "eventType":"SHUTDOWN", "record":{}}`,
+			collectorError: fmt.Errorf("test shutdown error"),
+			err:            fmt.Errorf("test shutdown error"),
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			logger := zaptest.NewLogger(t)
+			var ctx context.Context
+
+			if tc.cancel {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithCancel(context.Background())
+				cancel()
+			} else {
+				ctx = context.Background()
+			}
+
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(200)
+				w.Write([]byte(tc.serverResponse))
+				_, err := io.ReadAll(r.Body)
+				require.NoError(t, err, "failed to read request body: %v", err)
+			}))
+			defer server.Close()
+			u, err := url.Parse(server.URL)
+			require.NoError(t, err)
+
+			lm := manager{
+				collector:       &MockCollector{err: tc.collectorError},
+				logger:          logger,
+				listener:        telemetryapi.NewListener(logger),
+				extensionClient: extensionapi.NewClient(logger, fmt.Sprintf("%s", u.Host)),
+			}
+			if tc.err != nil {
+				err = lm.processEvents(ctx)
+				require.Error(t, err)
+				require.ErrorContains(t, err, tc.err.Error())
+			} else {
+				require.NoError(t, lm.processEvents(ctx))
+			}
+
+		})
+	}
+
+}