PR suggestions

DataDog · Dec 18, 2023 · 099a6c1 · 099a6c1
1 parent 7d8abfd
commit 099a6c1
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 15 deletions.
diff --git a/comp/logs/agent/component.go b/comp/logs/agent/component.go
@@ -48,6 +48,9 @@ type ServerlessLogsAgent interface {
 // Mock implements mock-specific methods.
 type Mock interface {
 	Component
+
+	SetSources(sources *sources.LogSources)
+	AsComponent() Component
 }
 
 // Module defines the fx options for this component.

diff --git a/comp/logs/agent/mock.go b/comp/logs/agent/mock.go
@@ -9,12 +9,13 @@ import (
 	"context"
 	"time"
 
+	"go.uber.org/fx"
+
 	"github.com/DataDog/datadog-agent/pkg/logs/diagnostic"
 	"github.com/DataDog/datadog-agent/pkg/logs/pipeline"
 	"github.com/DataDog/datadog-agent/pkg/logs/schedulers"
 	"github.com/DataDog/datadog-agent/pkg/logs/sources"
 	"github.com/DataDog/datadog-agent/pkg/util/optional"
-	"go.uber.org/fx"
 )
 
 type mockLogsAgent struct {
@@ -39,18 +40,6 @@ func newMock(deps dependencies) optional.Option[Mock] {
 	return optional.NewOption[Mock](logsAgent)
 }
 
-// NewMock can be used in other packages using the log agent as a dependency.
-func NewMock(logSources *sources.LogSources) optional.Option[Component] {
-	logsAgent := &mockLogsAgent{
-		hasFlushed:      false,
-		addedSchedulers: make([]schedulers.Scheduler, 0),
-		isRunning:       false,
-		flushDelay:      0,
-		logSources:      logSources,
-	}
-	return optional.NewOption[Component](logsAgent)
-}
-
 func (a *mockLogsAgent) start(context.Context) error {
 	a.isRunning = true
 	return nil
@@ -73,6 +62,10 @@ func (a *mockLogsAgent) GetMessageReceiver() *diagnostic.BufferedMessageReceiver
 	return nil
 }
 
+func (a *mockLogsAgent) SetSources(sources *sources.LogSources) {
+	a.logSources = sources
+}
+
 func (a *mockLogsAgent) GetSources() *sources.LogSources {
 	return a.logSources
 }
@@ -98,3 +91,7 @@ func (a *mockLogsAgent) Flush(ctx context.Context) {
 func (a *mockLogsAgent) GetPipelineProvider() pipeline.Provider {
 	return nil
 }
+
+func (a *mockLogsAgent) AsComponent() Component {
+	return a
+}
diff --git a/comp/metadata/inventorychecks/inventorychecksimpl/inventorychecks_test.go b/comp/metadata/inventorychecks/inventorychecksimpl/inventorychecks_test.go
@@ -12,10 +12,12 @@ import (
 	"github.com/stretchr/testify/mock"
 	"go.uber.org/fx"
 
+	"github.com/DataDog/datadog-agent/comp/core"
 	"github.com/DataDog/datadog-agent/comp/core/config"
 	"github.com/DataDog/datadog-agent/comp/core/log/logimpl"
 	logagent "github.com/DataDog/datadog-agent/comp/logs/agent"
 	logConfig "github.com/DataDog/datadog-agent/comp/logs/agent/config"
+	"github.com/DataDog/datadog-agent/comp/metadata/inventoryagent"
 	"github.com/DataDog/datadog-agent/pkg/collector"
 	"github.com/DataDog/datadog-agent/pkg/collector/check"
 	checkid "github.com/DataDog/datadog-agent/pkg/collector/check/id"
@@ -137,11 +139,13 @@ func TestGetPayload(t *testing.T) {
 		Source:     "redis",
 	})
 	logSources.AddSource(src)
-	mockLogAgent := logagent.NewMock(logSources)
+	mockLogAgent := fxutil.Test[optional.Option[logagent.Mock]](t, logagent.MockModule(), core.MockBundle(), inventoryagent.MockModule())
+	logsAgent, _ := mockLogAgent.Get()
+	logsAgent.SetSources(logSources)
 
 	ic := getTestInventoryChecks(t,
 		optional.NewOption[collector.Collector](mockColl),
-		mockLogAgent,
+		optional.NewOption[logagent.Component](logsAgent.AsComponent()),
 		overrides,
 	)