From cd4d15ffc1f8dc23fa889d300e431c562b6ee396 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Edu=20G=C3=B3mez=20Escandell?=
 <edu.gomez.escandell@canonical.com>
Date: Thu, 8 Feb 2024 11:39:26 +0100
Subject: [PATCH] Add Stop to Landscape test's mock config

---
 .../proservices/landscape/executor_test.go    |  7 +-
 .../proservices/landscape/landscape_test.go   | 86 +++++++++++++------
 2 files changed, 65 insertions(+), 28 deletions(-)

diff --git a/windows-agent/internal/proservices/landscape/executor_test.go b/windows-agent/internal/proservices/landscape/executor_test.go
index 4f70e2927..9833a861d 100644
--- a/windows-agent/internal/proservices/landscape/executor_test.go
+++ b/windows-agent/internal/proservices/landscape/executor_test.go
@@ -455,10 +455,9 @@ func testReceiveCommand(t *testing.T, distrosettings distroSettings, testSetup f
 
 	// Set up agent components (config, database, etc.)
 	if tb.conf == nil {
-		tb.conf = &mockConfig{
-			proToken:              "TOKEN",
-			landscapeClientConfig: executeLandscapeConfigTemplate(t, defaultLandscapeConfig, "", lis.Addr()),
-		}
+		tb.conf = newMockConfig(ctx)
+		tb.conf.proToken = "TOKEN"
+		tb.conf.landscapeClientConfig = executeLandscapeConfigTemplate(t, defaultLandscapeConfig, "", lis.Addr())
 	}
 
 	db, err := database.New(ctx, t.TempDir(), tb.conf)
diff --git a/windows-agent/internal/proservices/landscape/landscape_test.go b/windows-agent/internal/proservices/landscape/landscape_test.go
index d015159b9..2ee89e602 100644
--- a/windows-agent/internal/proservices/landscape/landscape_test.go
+++ b/windows-agent/internal/proservices/landscape/landscape_test.go
@@ -118,21 +118,22 @@ func TestConnect(t *testing.T) {
 			lis, server, mockService := setUpLandscapeMock(t, ctx, "localhost:", p)
 			defer lis.Close()
 
-			conf := &mockConfig{
-				proToken: "TOKEN",
+			conf := newMockConfig(ctx)
+			defer conf.Stop()
 
-				// We trigger an error on first-contact SendUpdatedInfo by erroring out in conf.ProToken()
-				proTokenErr: tc.tokenErr,
+			conf.proToken = "TOKEN"
 
-				// We trigger errors trying to read or write to/from the registry
-				landscapeUIDErr:    tc.landscapeUIDReadErr,
-				setLandscapeUIDErr: tc.landscapeUIDWriteErr,
+			// We trigger an error on first-contact SendUpdatedInfo by erroring out in conf.ProToken()
+			conf.proTokenErr = tc.tokenErr
 
-				// We trigger an error when deciding to use a certificate or not
-				landscapeConfigErr: tc.breakLandscapeClientConfig,
+			// We trigger errors trying to read or write to/from the registry
+			conf.landscapeUIDErr = tc.landscapeUIDReadErr
+			conf.setLandscapeUIDErr = tc.landscapeUIDWriteErr
 
-				landscapeAgentUID: tc.uid,
-			}
+			// We trigger an error when deciding to use a certificate or not
+			conf.landscapeConfigErr = tc.breakLandscapeClientConfig
+
+			conf.landscapeAgentUID = tc.uid
 
 			if tc.emptyToken {
 				conf.proToken = ""
@@ -244,10 +245,11 @@ func TestSendUpdatedInfo(t *testing.T) {
 
 			lis, server, mockService := setUpLandscapeMock(t, ctx, "localhost:", "")
 
-			conf := &mockConfig{
-				proToken:              "TOKEN",
-				landscapeClientConfig: executeLandscapeConfigTemplate(t, defaultLandscapeConfig, "", lis.Addr()),
-			}
+			conf := newMockConfig(ctx)
+			defer conf.Stop()
+
+			conf.proToken = "TOKEN"
+			conf.landscapeClientConfig = executeLandscapeConfigTemplate(t, defaultLandscapeConfig, "", lis.Addr())
 
 			//nolint:errcheck // We don't care about these errors
 			go server.Serve(lis)
@@ -439,10 +441,11 @@ func TestAutoReconnection(t *testing.T) {
 			defer lis.Close()
 			defer server.Stop()
 
-			conf := &mockConfig{
-				proToken:              "TOKEN",
-				landscapeClientConfig: executeLandscapeConfigTemplate(t, defaultLandscapeConfig, "", lis.Addr()),
-			}
+			conf := newMockConfig(ctx)
+			defer conf.Stop()
+
+			conf.proToken = "TOKEN"
+			conf.landscapeClientConfig = executeLandscapeConfigTemplate(t, defaultLandscapeConfig, "", lis.Addr())
 
 			db, err := database.New(ctx, t.TempDir(), conf)
 			require.NoError(t, err, "Setup: database New should not return an error")
@@ -610,10 +613,11 @@ func TestReconnect(t *testing.T) {
 
 			lis, server, _ := setUpLandscapeMock(t, ctx, "localhost:", certPath)
 
-			conf := &mockConfig{
-				proToken:              "TOKEN",
-				landscapeClientConfig: executeLandscapeConfigTemplate(t, lcapeConfig, certPath, lis.Addr()),
-			}
+			conf := newMockConfig(ctx)
+			defer conf.Stop()
+
+			conf.proToken = "TOKEN"
+			conf.landscapeClientConfig = executeLandscapeConfigTemplate(t, lcapeConfig, certPath, lis.Addr())
 
 			//nolint:errcheck // We don't care about these errors
 			go server.Serve(lis)
@@ -737,6 +741,9 @@ func setUpLandscapeMock(t *testing.T, ctx context.Context, addr string, certPath
 }
 
 type mockConfig struct {
+	ctx    context.Context
+	cancel func()
+
 	proToken              string
 	landscapeClientConfig string
 	landscapeAgentUID     string
@@ -747,10 +754,29 @@ type mockConfig struct {
 	setLandscapeUIDErr bool
 
 	callbacks []func()
+	wg        sync.WaitGroup
 
 	mu sync.Mutex
 }
 
+func newMockConfig(ctx context.Context) *mockConfig {
+	ctx, cancel := context.WithCancel(ctx)
+
+	return &mockConfig{
+		ctx:    ctx,
+		cancel: cancel,
+	}
+}
+
+func (m *mockConfig) Stop() {
+	m.cancel()
+
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	m.wg.Wait()
+}
+
 func (m *mockConfig) LandscapeClientConfig(ctx context.Context) (string, config.Source, error) {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -811,6 +837,18 @@ func (m *mockConfig) Notify(f func()) {
 
 func (m *mockConfig) triggerNotifications() {
 	for _, f := range m.callbacks {
-		go f()
+		m.wg.Add(1)
+		f := f
+		go func() {
+			defer m.wg.Done()
+
+			select {
+			case <-m.ctx.Done():
+				return
+			default:
+			}
+
+			f()
+		}()
 	}
 }