[cleanup] Use storagetest.NewStorageHost() from contrib (#5744)

## Which problem is this PR solving? - Simplify code Signed-off-by: Yuri Shkuro <[email protected]>
jaegertracing · Jul 14, 2024 · 1e1b0a1 · 1e1b0a1
1 parent 0616fba
commit 1e1b0a1
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 73 deletions.
diff --git a/cmd/jaeger/internal/exporters/storageexporter/exporter_test.go b/cmd/jaeger/internal/exporters/storageexporter/exporter_test.go
@@ -124,8 +124,7 @@ func TestExporter(t *testing.T) {
 
 	host := makeStorageExtension(t, memstoreName)
 
-	err = tracesExporter.Start(ctx, host)
-	require.NoError(t, err)
+	require.NoError(t, tracesExporter.Start(ctx, host))
 	defer func() {
 		require.NoError(t, tracesExporter.Shutdown(ctx))
 	}()

diff --git a/cmd/jaeger/internal/extension/jaegerquery/server_test.go b/cmd/jaeger/internal/extension/jaegerquery/server_test.go
@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/storage/storagetest"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/collector/component"
@@ -80,27 +81,6 @@ func (fakeStorageExt) Shutdown(context.Context) error {
 	return nil
 }
 
-type storageHost struct {
-	extension component.Component
-}
-
-func (storageHost) ReportFatalError(error) {
-}
-
-func (host storageHost) GetExtensions() map[component.ID]component.Component {
-	return map[component.ID]component.Component{
-		jaegerstorage.ID: host.extension,
-	}
-}
-
-func (storageHost) GetFactory(_ component.Kind, _ component.Type) component.Factory {
-	return nil
-}
-
-func (storageHost) GetExporters() map[component.DataType]map[component.ID]component.Component {
-	return nil
-}
-
 func TestServerDependencies(t *testing.T) {
 	expectedDependencies := []component.ID{jaegerstorage.ID}
 	telemetrySettings := component.TelemetrySettings{
@@ -114,9 +94,7 @@ func TestServerDependencies(t *testing.T) {
 }
 
 func TestServerStart(t *testing.T) {
-	host := storageHost{
-		extension: fakeStorageExt{},
-	}
+	host := storagetest.NewStorageHost().WithExtension(jaegerstorage.ID, fakeStorageExt{})
 	tests := []struct {
 		name        string
 		config      *Config
@@ -251,7 +229,7 @@ func TestServerAddArchiveStorage(t *testing.T) {
 			}
 			server := newServer(tt.config, telemetrySettings)
 			if tt.extension != nil {
-				host = storageHost{extension: tt.extension}
+				host = storagetest.NewStorageHost().WithExtension(jaegerstorage.ID, tt.extension)
 			}
 			err := server.addArchiveStorage(tt.qSvcOpts, host)
 			if tt.expectedErr == "" {

diff --git a/cmd/jaeger/internal/extension/jaegerstorage/extension_test.go b/cmd/jaeger/internal/extension/jaegerstorage/extension_test.go
@@ -11,6 +11,7 @@ import (
 	"net/http/httptest"
 	"testing"
 
+	"github.com/open-telemetry/opentelemetry-collector-contrib/extension/storage/storagetest"
 	"github.com/stretchr/testify/require"
 	"go.opentelemetry.io/collector/component"
 	"go.opentelemetry.io/collector/component/componenttest"
@@ -31,26 +32,6 @@ import (
 	"github.com/jaegertracing/jaeger/storage/spanstore"
 )
 
-type storageHost struct {
-	t   *testing.T
-	ext component.Component
-}
-
-func (host storageHost) GetExtensions() map[component.ID]component.Component {
-	return map[component.ID]component.Component{
-		ID: host.ext,
-	}
-}
-
-func (host storageHost) ReportFatalError(err error) {
-	host.t.Fatal(err)
-}
-
-func (storageHost) GetFactory(_ component.Kind, _ component.Type) component.Factory { return nil }
-func (storageHost) GetExporters() map[component.DataType]map[component.ID]component.Component {
-	return nil
-}
-
 type errorFactory struct {
 	closeErr error
 }
@@ -81,7 +62,7 @@ func TestStorageFactoryBadHostError(t *testing.T) {
 }
 
 func TestStorageFactoryBadNameError(t *testing.T) {
-	host := storageHost{t: t, ext: startStorageExtension(t, "foo")}
+	host := storagetest.NewStorageHost().WithExtension(ID, startStorageExtension(t, "foo"))
 	_, err := GetStorageFactory("bar", host)
 	require.ErrorContains(t, err, "cannot find definition of storage 'bar'")
 }
@@ -105,7 +86,7 @@ func TestGetFactoryV2Error(t *testing.T) {
 
 func TestGetFactory(t *testing.T) {
 	const name = "foo"
-	host := storageHost{t: t, ext: startStorageExtension(t, name)}
+	host := storagetest.NewStorageHost().WithExtension(ID, startStorageExtension(t, name))
 	f, err := GetStorageFactory(name, host)
 	require.NoError(t, err)
 	require.NotNil(t, f)

diff --git a/cmd/jaeger/internal/processors/adaptivesampling/processor_test.go b/cmd/jaeger/internal/processors/adaptivesampling/processor_test.go
@@ -25,26 +25,6 @@ import (
 	"github.com/jaegertracing/jaeger/plugin/storage/memory"
 )
 
-type samplingHost struct {
-	t                 *testing.T
-	samplingExtension component.Component
-}
-
-func (host samplingHost) GetExtensions() map[component.ID]component.Component {
-	return map[component.ID]component.Component{
-		remotesampling.ID: host.samplingExtension,
-	}
-}
-
-func (host samplingHost) ReportFatalError(err error) {
-	host.t.Fatal(err)
-}
-
-func (samplingHost) GetFactory(_ component.Kind, _ component.Type) component.Factory { return nil }
-func (samplingHost) GetExporters() map[component.DataType]map[component.ID]component.Component {
-	return nil
-}
-
 func makeStorageExtension(t *testing.T, memstoreName string) component.Host {
 	telemetrySettings := component.TelemetrySettings{
 		Logger:         zaptest.NewLogger(t),
@@ -74,7 +54,7 @@ func makeStorageExtension(t *testing.T, memstoreName string) component.Host {
 
 var _ component.Config = (*Config)(nil)
 
-func makeRemoteSamplingExtension(t *testing.T, cfg component.Config) samplingHost {
+func makeRemoteSamplingExtension(t *testing.T, cfg component.Config) component.Host {
 	extensionFactory := remotesampling.NewFactory()
 	samplingExtension, err := extensionFactory.CreateExtension(
 		context.Background(),
@@ -87,7 +67,7 @@ func makeRemoteSamplingExtension(t *testing.T, cfg component.Config) samplingHos
 		cfg,
 	)
 	require.NoError(t, err)
-	host := samplingHost{t: t, samplingExtension: samplingExtension}
+	host := storagetest.NewStorageHost().WithExtension(remotesampling.ID, samplingExtension)
 	storageHost := makeStorageExtension(t, "foobar")
 
 	err = samplingExtension.Start(context.Background(), storageHost)
@@ -148,9 +128,8 @@ func makeTracesOneSpan() ptrace.Traces {
 }
 
 func TestGetAdaptiveSamplingComponentsError(t *testing.T) {
-	host := &samplingHost{}
 	processor := &traceProcessor{}
-	err := processor.start(context.Background(), host)
+	err := processor.start(context.Background(), storagetest.NewStorageHost())
 	require.ErrorContains(t, err, "cannot load adaptive sampling components")
 }