Add tests for internal/extension/jaegerquery (#5123)

## Which problem is this PR solving?
- Part of #5068

## Description of the changes
- This commit adds tests for the
`cmd/jaeger/internal/extension/jaegerquery` package.

## How was this change tested?
- make test

## Checklist
- [x] I have read
https://github.com/jaegertracing/jaeger/blob/master/CONTRIBUTING_GUIDELINES.md
- [x] I have signed all commits
- [x] I have added unit tests for the new functionality
- [x] I have run lint and test steps successfully
  - for `jaeger`: `make lint test`
  - for `jaeger-ui`: `yarn lint` and `yarn test`

---------

Signed-off-by: VaibhavMalik4187 <[email protected]>
Signed-off-by: Yuri Shkuro <[email protected]>
Signed-off-by: Yuri Shkuro <[email protected]>
Co-authored-by: Yuri Shkuro <[email protected]>
Co-authored-by: Yuri Shkuro <[email protected]>

cmd/jaeger/internal/extension/jaegerquery/config_test.go

@@ -0,0 +1,43 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerquery
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_Validate(t *testing.T) {
+	tests := []struct {
+		name        string
+		config      *Config
+		expectedErr string
+	}{
+		{
+			name:        "Empty config",
+			config:      &Config{},
+			expectedErr: "TraceStoragePrimary: non zero value required",
+		},
+		{
+			name: "Non empty-config",
+			config: &Config{
+				TraceStoragePrimary: "some-storage",
+			},
+			expectedErr: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := tt.config.Validate()
+			if tt.expectedErr == "" {
+				require.NoError(t, err)
+			} else {
+				assert.Equal(t, tt.expectedErr, err.Error())
+			}
+		})
+	}
+}

cmd/jaeger/internal/extension/jaegerquery/factory_test.go

@@ -0,0 +1,30 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerquery
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/extension"
+)
+
+func Test_NewFactory(t *testing.T) {
+	factory := NewFactory()
+	assert.Equal(t, componentType, factory.Type())
+	assert.Equal(t, factory.CreateDefaultConfig(), createDefaultConfig())
+}
+
+func Test_CreateExtension(t *testing.T) {
+	set := extension.CreateSettings{
+		ID: ID,
+	}
+	cfg := createDefaultConfig()
+	extension, err := createExtension(context.Background(), set, cfg)
+
+	require.NoError(t, err)
+	assert.NotNil(t, extension)
+}

cmd/jaeger/internal/extension/jaegerquery/package_test.go

@@ -0,0 +1,14 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerquery
+
+import (
+	"testing"
+
+	"github.com/jaegertracing/jaeger/pkg/testutils"
+)
+
+func TestMain(m *testing.M) {
+	testutils.VerifyGoLeaks(m)
+}

cmd/jaeger/internal/extension/jaegerquery/server.go

@@ -5,6 +5,7 @@ package jaegerquery
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"go.opentelemetry.io/collector/component"
@@ -27,9 +28,10 @@ var (
 )
 
 type server struct {
-	config *Config
-	logger *zap.Logger
-	server *queryApp.Server
+	config  *Config
+	logger  *zap.Logger
+	server  *queryApp.Server
+	jtracer *jtracer.JTracer
 }
 
 func newServer(config *Config, otel component.TelemetrySettings) *server {
@@ -73,7 +75,7 @@ func (s *server) Start(ctx context.Context, host component.Host) error {
 	// TODO OTel-collector does not initialize the tracer currently
 	// https://github.com/open-telemetry/opentelemetry-collector/issues/7532
 	//nolint
-	jtracer, err := jtracer.New("jaeger")
+	s.jtracer, err = jtracer.New("jaeger")
 	if err != nil {
 		return fmt.Errorf("could not initialize a tracer: %w", err)
 	}
@@ -88,7 +90,7 @@ func (s *server) Start(ctx context.Context, host component.Host) error {
 		metricsQueryService,
 		s.makeQueryOptions(),
 		tm,
-		jtracer,
+		s.jtracer,
 	)
 	if err != nil {
 		return fmt.Errorf("could not create jaeger-query: %w", err)
@@ -129,8 +131,12 @@ func (s *server) makeQueryOptions() *queryApp.QueryOptions {
 }
 
 func (s *server) Shutdown(ctx context.Context) error {
+	var errs []error
 	if s.server != nil {
-		return s.server.Close()
+		errs = append(errs, s.server.Close())
 	}
-	return nil
+	if s.jtracer != nil {
+		errs = append(errs, s.jtracer.Close(ctx))
+	}
+	return errors.Join(errs...)
 }

cmd/jaeger/internal/extension/jaegerquery/server_test.go

@@ -0,0 +1,237 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package jaegerquery
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/component/componenttest"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zaptest"
+
+	"github.com/jaegertracing/jaeger/cmd/jaeger/internal/extension/jaegerstorage"
+	"github.com/jaegertracing/jaeger/cmd/query/app/querysvc"
+	"github.com/jaegertracing/jaeger/pkg/metrics"
+	"github.com/jaegertracing/jaeger/pkg/testutils"
+	"github.com/jaegertracing/jaeger/storage"
+	"github.com/jaegertracing/jaeger/storage/dependencystore"
+	depsmocks "github.com/jaegertracing/jaeger/storage/dependencystore/mocks"
+	"github.com/jaegertracing/jaeger/storage/spanstore"
+	spanstoremocks "github.com/jaegertracing/jaeger/storage/spanstore/mocks"
+)
+
+type fakeFactory struct {
+	name string
+}
+
+func (ff fakeFactory) CreateDependencyReader() (dependencystore.Reader, error) {
+	if ff.name == "need-dependency-reader-error" {
+		return nil, fmt.Errorf("test-error")
+	}
+	return &depsmocks.Reader{}, nil
+}
+
+func (ff fakeFactory) CreateSpanReader() (spanstore.Reader, error) {
+	if ff.name == "need-span-reader-error" {
+		return nil, fmt.Errorf("test-error")
+	}
+	return &spanstoremocks.Reader{}, nil
+}
+
+func (ff fakeFactory) CreateSpanWriter() (spanstore.Writer, error) {
+	if ff.name == "need-span-writer-error" {
+		return nil, fmt.Errorf("test-error")
+	}
+	return &spanstoremocks.Writer{}, nil
+}
+
+func (ff fakeFactory) Initialize(metrics.Factory, *zap.Logger) error {
+	if ff.name == "need-initialize-error" {
+		return fmt.Errorf("test-error")
+	}
+	return nil
+}
+
+type fakeStorageExt struct{}
+
+var _ jaegerstorage.Extension = (*fakeStorageExt)(nil)
+
+func (fse fakeStorageExt) Factory(name string) (storage.Factory, bool) {
+	if name == "need-factory-error" {
+		return nil, false
+	}
+	return fakeFactory{name: name}, true
+}
+
+func (fse fakeStorageExt) Start(ctx context.Context, host component.Host) error {
+	return nil
+}
+
+func (fse fakeStorageExt) Shutdown(ctx context.Context) error {
+	return nil
+}
+
+type storageHost struct {
+	extension component.Component
+}
+
+func (host storageHost) ReportFatalError(err 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{
+		Logger: zaptest.NewLogger(t),
+	}
+
+	server := newServer(createDefaultConfig().(*Config), telemetrySettings)
+	dependencies := server.Dependencies()
+
+	assert.Equal(t, expectedDependencies, dependencies)
+}
+
+func TestServerStart(t *testing.T) {
+	host := storageHost{
+		extension: fakeStorageExt{},
+	}
+	tests := []struct {
+		name        string
+		config      *Config
+		expectedErr string
+	}{
+		{
+			name: "Non-empty config with fake storage host",
+			config: &Config{
+				TraceStorageArchive: "jaeger_storage",
+				TraceStoragePrimary: "jaeger_storage",
+			},
+		},
+		{
+			name: "factory error",
+			config: &Config{
+				TraceStoragePrimary: "need-factory-error",
+			},
+			expectedErr: "cannot find primary storage",
+		},
+		{
+			name: "span reader error",
+			config: &Config{
+				TraceStoragePrimary: "need-span-reader-error",
+			},
+			expectedErr: "cannot create span reader",
+		},
+		{
+			name: "dependency error",
+			config: &Config{
+				TraceStoragePrimary: "need-dependency-reader-error",
+			},
+			expectedErr: "cannot create dependencies reader",
+		},
+		{
+			name: "storage archive error",
+			config: &Config{
+				TraceStorageArchive: "need-factory-error",
+				TraceStoragePrimary: "jaeger_storage",
+			},
+			expectedErr: "cannot find archive storage factory",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			telemetrySettings := component.TelemetrySettings{
+				Logger: zaptest.NewLogger(t),
+			}
+			server := newServer(tt.config, telemetrySettings)
+			err := server.Start(context.Background(), host)
+
+			if tt.expectedErr == "" {
+				require.NoError(t, err)
+				defer server.Shutdown(context.Background())
+			} else {
+				require.ErrorContains(t, err, tt.expectedErr)
+			}
+		})
+	}
+}
+
+func TestServerAddArchiveStorage(t *testing.T) {
+	host := componenttest.NewNopHost()
+
+	tests := []struct {
+		name           string
+		qSvcOpts       *querysvc.QueryServiceOptions
+		config         *Config
+		extension      component.Component
+		expectedOutput string
+		expectedErr    string
+	}{
+		{
+			name:           "Archive storage unset",
+			config:         &Config{},
+			qSvcOpts:       &querysvc.QueryServiceOptions{},
+			expectedOutput: `{"level":"info","msg":"Archive storage not configured"}` + "\n",
+			expectedErr:    "",
+		},
+		{
+			name: "Archive storage set",
+			config: &Config{
+				TraceStorageArchive: "random-value",
+			},
+			qSvcOpts:       &querysvc.QueryServiceOptions{},
+			expectedOutput: "",
+			expectedErr:    "cannot find archive storage factory: cannot find extension",
+		},
+		{
+			name: "Archive storage not supported",
+			config: &Config{
+				TraceStorageArchive: "badger",
+			},
+			qSvcOpts:       &querysvc.QueryServiceOptions{},
+			extension:      fakeStorageExt{},
+			expectedOutput: "Archive storage not supported by the factory",
+			expectedErr:    "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			logger, buf := testutils.NewLogger()
+			telemetrySettings := component.TelemetrySettings{
+				Logger: logger,
+			}
+			server := newServer(tt.config, telemetrySettings)
+			if tt.extension != nil {
+				host = storageHost{extension: tt.extension}
+			}
+			err := server.addArchiveStorage(tt.qSvcOpts, host)
+			if tt.expectedErr == "" {
+				require.NoError(t, err)
+			} else {
+				require.ErrorContains(t, err, tt.expectedErr)
+			}
+
+			assert.Contains(t, buf.String(), tt.expectedOutput)
+		})
+	}
+}

cmd/jaeger/internal/extension/jaegerstorage/extension.go

@@ -57,7 +57,7 @@ func GetStorageFactory(name string, host component.Host) (storage.Factory, error
 	f, ok := comp.(Extension).Factory(name)
 	if !ok {
 		return nil, fmt.Errorf(
-			"cannot find storage '%s' declared with '%s' extension",
+			"cannot find storage '%s' declared by '%s' extension",
 			name, componentType,
 		)
 	}