Fix configFromUnit test, add logging for e2e tests (#5396) (#5405)

* Fix configFromUnit test, add logging for e2e tests

* Fix linter

* Fix data race

* linter suggestions

* Apply suggestions from code review



---------


(cherry picked from commit 35012a1)

Co-authored-by: Michel Laterman <[email protected]>
Co-authored-by: Shaunak Kashyap <[email protected]>

Author: 3 people

internal/pkg/server/agent.go

@@ -282,50 +282,48 @@ func (a *Agent) unitModified(ctx context.Context, unit *client.Unit) error {
 			// not our input unit; would have been marked failed in unitAdded; do nothing
 			return nil
 		}
-		if exp.State == client.UnitStateHealthy {
+		switch exp.State {
+		case client.UnitStateHealthy:
 			if a.outputUnit == nil {
 				// still no output unit; would have been marked starting already; do nothing
 				return nil
 			}
-
 			// configuration modified (should still be running)
 			return a.reconfigure(ctx)
-		} else if exp.State == client.UnitStateStopped {
+		case client.UnitStateStopped:
 			// unit should be stopped
 			a.stop()
 			return nil
+		default:
+			return fmt.Errorf("unknown unit state %v", exp.State)
 		}
-		return fmt.Errorf("unknown unit state %v", exp.State)
 	}
 	if unit.Type() == client.UnitTypeOutput {
 		if a.outputUnit != unit {
 			// not our output unit; would have been marked failed in unitAdded; do nothing
 			return nil
 		}
-		if exp.State == client.UnitStateHealthy {
+		switch exp.State {
+		case client.UnitStateHealthy:
 			if a.inputUnit == nil {
 				// still no input unit; would have been marked starting already; do nothing
 				return nil
 			}
-
 			// configuration modified (should still be running)
 			return a.reconfigure(ctx)
-		} else if exp.State == client.UnitStateStopped {
+		case client.UnitStateStopped:
 			// unit should be stopped
 			a.stop()
 			return nil
+		default:
+			return fmt.Errorf("unknown unit state %v", exp.State)
 		}
-		return fmt.Errorf("unknown unit state %v", exp.State)
 	}
 	return fmt.Errorf("unknown unit type %v", unit.Type())
 }
 
 func (a *Agent) unitRemoved(unit *client.Unit) {
-	stop := false
 	if a.inputUnit == unit || a.outputUnit == unit {
-		stop = true
-	}
-	if stop {
 		a.stop()
 	}
 	if a.inputUnit == unit {
@@ -726,8 +724,7 @@ func (a *Agent) esOutputCheck(ctx context.Context, cfg map[string]interface{}) e
 func (a *Agent) esOutputCheckLoop(ctx context.Context, delay time.Duration, cfg map[string]interface{}) {
 	for {
 		if err := sleep.WithContext(ctx, delay); err != nil {
-			zerolog.Ctx(ctx).Debug().Msg("Async output check context cancelled")
-			return
+			return // context cancelled
 		}
 		err := a.esOutputCheck(ctx, cfg)
 		if err == nil {

internal/pkg/server/agent_test.go

@@ -11,16 +11,17 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/types/known/structpb"
+
 	"github.com/elastic/elastic-agent-client/v7/pkg/client"
 	"github.com/elastic/elastic-agent-client/v7/pkg/proto"
 	"github.com/elastic/fleet-server/v7/internal/pkg/build"
 	testlog "github.com/elastic/fleet-server/v7/internal/pkg/testing/log"
 	"github.com/elastic/fleet-server/v7/version"
 	"github.com/elastic/go-ucfg"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/mock"
-	"github.com/stretchr/testify/require"
-	"google.golang.org/protobuf/types/known/structpb"
 )
 
 func TestCLIOverrides(t *testing.T) {
@@ -88,7 +89,7 @@ func TestCLIOverrides(t *testing.T) {
 		agent:      clientMock,
 	}
 
-	generatedCfg, err := agent.configFromUnits(context.Background())
+	generatedCfg, err := agent.configFromUnits(t.Context())
 	require.NoError(t, err)
 	require.Equal(t, httpEnabledExpected, generatedCfg.HTTP.Enabled)
 	require.Equal(t, httpHostExpected, generatedCfg.HTTP.Host)
@@ -108,29 +109,29 @@ func (c *mockClientV2) RegisterOptionalDiagnosticHook(paramTag string, name stri
 
 func (c *mockClientV2) Start(ctx context.Context) error {
 	args := c.Called()
-	return args.Get(0).(error)
+	return args.Get(0).(error) //nolint:errcheck // testing mock
 }
 
 func (c *mockClientV2) Stop() {}
 
 func (c *mockClientV2) UnitChanges() <-chan client.UnitChanged {
 	args := c.Called()
-	return args.Get(0).(<-chan client.UnitChanged)
+	return args.Get(0).(<-chan client.UnitChanged) //nolint:errcheck // testing mock
 }
 
 func (c *mockClientV2) Errors() <-chan error {
 	args := c.Called()
-	return args.Get(0).(<-chan error)
+	return args.Get(0).(<-chan error) //nolint:errcheck // testing mock
 }
 
 func (c *mockClientV2) Artifacts() client.ArtifactsClient {
 	args := c.Called()
-	return args.Get(0).(client.ArtifactsClient)
+	return args.Get(0).(client.ArtifactsClient) //nolint:errcheck // testing mock
 }
 
 func (c *mockClientV2) AgentInfo() *client.AgentInfo {
 	args := c.Called()
-	return args.Get(0).(*client.AgentInfo)
+	return args.Get(0).(*client.AgentInfo) //nolint:errcheck // testing mock
 }
 
 type mockClientUnit struct {
@@ -140,12 +141,12 @@ type mockClientUnit struct {
 func (u *mockClientUnit) Expected() client.Expected {
 	args := u.Called()
 
-	return args.Get(0).(client.Expected)
+	return args.Get(0).(client.Expected) //nolint:errcheck // testing mock
 }
 
 func (u *mockClientUnit) UpdateState(state client.UnitState, message string, payload map[string]interface{}) error {
 	args := u.Called()
-	return args.Get(0).(error)
+	return args.Get(0).(error) //nolint:errcheck // testing mock
 }
 
 func Test_Agent_configFromUnits(t *testing.T) {
@@ -194,7 +195,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			outputUnit: mockOutClient,
 		}
 
-		cfg, err := a.configFromUnits(context.Background())
+		cfg, err := a.configFromUnits(t.Context())
 		require.NoError(t, err)
 		require.Len(t, cfg.Inputs, 1)
 		assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
@@ -234,7 +235,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			outputUnit: mockOutClient,
 		}
 
-		cfg, err := a.configFromUnits(context.Background())
+		cfg, err := a.configFromUnits(t.Context())
 		require.NoError(t, err)
 		assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
 		require.Len(t, cfg.Output.Elasticsearch.Hosts, 2)
@@ -295,7 +296,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			outputUnit: mockOutClient,
 		}
 
-		cfg, err := a.configFromUnits(context.Background())
+		cfg, err := a.configFromUnits(t.Context())
 		require.NoError(t, err)
 		require.Len(t, cfg.Inputs, 1)
 		assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
@@ -375,7 +376,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			outputUnit: mockOutClient,
 		}
 
-		cfg, err := a.configFromUnits(context.Background())
+		cfg, err := a.configFromUnits(t.Context())
 		require.NoError(t, err)
 		require.Len(t, cfg.Inputs, 1)
 		assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
@@ -439,7 +440,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			outputUnit: mockOutClient,
 		}
 
-		cfg, err := a.configFromUnits(context.Background())
+		cfg, err := a.configFromUnits(t.Context())
 		require.NoError(t, err)
 		require.Len(t, cfg.Inputs, 1)
 		assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
@@ -518,7 +519,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			outputUnit: mockOutClient,
 		}
 
-		cfg, err := a.configFromUnits(context.Background())
+		cfg, err := a.configFromUnits(t.Context())
 		require.NoError(t, err)
 		require.Len(t, cfg.Inputs, 1)
 		assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
@@ -571,7 +572,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			outputUnit: mockOutClient,
 		}
 
-		cfg, err := a.configFromUnits(context.Background())
+		cfg, err := a.configFromUnits(t.Context())
 		require.NoError(t, err)
 		require.Len(t, cfg.Inputs, 1)
 		assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
@@ -625,7 +626,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			outputUnit: mockOutClient,
 		}
 
-		cfg, err := a.configFromUnits(context.Background())
+		cfg, err := a.configFromUnits(t.Context())
 		require.NoError(t, err)
 		require.Len(t, cfg.Inputs, 1)
 		assert.Equal(t, "fleet-server", cfg.Inputs[0].Type)
@@ -684,7 +685,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			},
 		}
 
-		ctx := testlog.SetLogger(t).WithContext(context.Background())
+		ctx := testlog.SetLogger(t).WithContext(t.Context())
 		cfg, err := a.configFromUnits(ctx)
 		require.NoError(t, err)
 		assert.Equal(t, "test-token", cfg.Output.Elasticsearch.ServiceToken)
@@ -734,7 +735,7 @@ func Test_Agent_configFromUnits(t *testing.T) {
 			},
 		}
 
-		ctx := testlog.SetLogger(t).WithContext(context.Background())
+		ctx := testlog.SetLogger(t).WithContext(t.Context())
 		cfg, err := a.configFromUnits(ctx)
 		require.NoError(t, err)
 		assert.Equal(t, "test-token", cfg.Output.Elasticsearch.ServiceToken)
@@ -839,7 +840,7 @@ func TestInjectMissingOutputAttributes(t *testing.T) {
 	}
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			injectMissingOutputAttributes(context.Background(), tc.input, bootstrap)
+			injectMissingOutputAttributes(t.Context(), tc.input, bootstrap)
 			assert.Equal(t, len(tc.expect), len(tc.input), "expected map sizes don't match")
 			assert.Equal(t, tc.expect, tc.input)
 		})
@@ -914,7 +915,7 @@ func TestInjectMissingOutputAttributes(t *testing.T) {
 
 	for _, tc := range sslTests {
 		t.Run(tc.name, func(t *testing.T) {
-			injectMissingOutputAttributes(context.Background(), tc.input, bootstrapVerifyNone)
+			injectMissingOutputAttributes(t.Context(), tc.input, bootstrapVerifyNone)
 			assert.Equal(t, len(tc.expect), len(tc.input), "expected map sizes don't match")
 			assert.Equal(t, tc.expect, tc.input)
 		})
@@ -928,7 +929,7 @@ func Test_Agent_esOutputCheckLoop(t *testing.T) {
 			chReconfigure: make(chan struct{}, 1),
 		}
 
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(t.Context())
 		cancel()
 		a.esOutputCheckLoop(ctx, time.Millisecond*10, map[string]interface{}{})
 		assert.Empty(t, a.chReconfigure)
@@ -948,7 +949,7 @@ func Test_Agent_esOutputCheckLoop(t *testing.T) {
 			},
 			chReconfigure: make(chan struct{}, 1),
 		}
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(t.Context())
 		defer cancel()
 		a.esOutputCheckLoop(ctx, time.Millisecond*10, map[string]interface{}{
 			"service_token": "test-token",
@@ -977,7 +978,7 @@ func Test_Agent_esOutputCheckLoop(t *testing.T) {
 			},
 			chReconfigure: make(chan struct{}, 1),
 		}
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(t.Context())
 		defer cancel()
 		a.esOutputCheckLoop(ctx, time.Millisecond*10, map[string]interface{}{
 			"service_token": "test-token",

internal/pkg/ver/check.go

@@ -13,9 +13,10 @@ import (
 	"strconv"
 	"strings"
 
-	esh "github.com/elastic/fleet-server/v7/internal/pkg/es"
 	"github.com/rs/zerolog"
 
+	esh "github.com/elastic/fleet-server/v7/internal/pkg/es"
+
 	"github.com/hashicorp/go-version"
 
 	"github.com/elastic/go-elasticsearch/v8"
@@ -30,15 +31,24 @@ var (
 // CheckCompatiblility will check the remote Elasticsearch version retrieved by the Elasticsearch client with the passed fleet version.
 // Versions are compatible when Elasticsearch's version is greater then or equal to fleet-server's version
 func CheckCompatibility(ctx context.Context, esCli *elasticsearch.Client, fleetVersion string) (string, error) {
-	zerolog.Ctx(ctx).Debug().Str("fleet_version", fleetVersion).Msg("check version compatibility with elasticsearch")
+	// Version checks may run concurrently with other operations
+	// This can cause some flakiness with tests so we need to get the logger from the context before its cancelled
+	var logger *zerolog.Logger
+	select {
+	case <-ctx.Done():
+		return "", ctx.Err()
+	default:
+		logger = zerolog.Ctx(ctx)
+	}
+	logger.Debug().Str("fleet_version", fleetVersion).Msg("check version compatibility with elasticsearch")
 
 	esVersion, err := esh.FetchESVersion(ctx, esCli)
 
 	if err != nil {
-		zerolog.Ctx(ctx).Error().Err(err).Msg("failed to fetch elasticsearch version")
+		logger.Error().Err(err).Msg("failed to fetch elasticsearch version")
 		return "", err
 	}
-	zerolog.Ctx(ctx).Debug().Str("elasticsearch_version", esVersion).Msg("fetched elasticsearch version")
+	logger.Debug().Str("elasticsearch_version", esVersion).Msg("fetched elasticsearch version")
 
 	return esVersion, checkCompatibility(ctx, fleetVersion, esVersion)
 }

testing/e2e/scaffold/scaffold.go

@@ -244,20 +244,25 @@ func (s *Scaffold) AgentIsOnline(ctx context.Context, id string) {
 			req.Header.Set("kbn-xsrf", "e2e-setup")
 
 			resp, err := s.Client.Do(req)
-			s.Require().NoError(err)
-			defer resp.Body.Close()
+			s.Require().NoError(err, "kibana agent request failure")
 			if resp.StatusCode != http.StatusOK {
 				timer.Reset(time.Second)
+				resp.Body.Close()
 				continue
 			}
 
+			p, err := io.ReadAll(resp.Body)
+			s.Require().NoError(err, "unable to read kibana agent response")
+			resp.Body.Close()
+
 			var obj struct {
 				Item struct {
 					Status string `json:"status"`
 				} `json:"item"`
 			}
-			err = json.NewDecoder(resp.Body).Decode(&obj)
-			s.Require().NoError(err)
+			s.T().Logf("Kibana agent response: %s", string(p))
+			err = json.Unmarshal(p, &obj)
+			s.Require().NoError(err, "unmarshal failure")
 			if obj.Item.Status == "online" {
 				return
 			}