mirror of https://github.com/grafana/grafana.git
343 lines
8.7 KiB
Go
343 lines
8.7 KiB
Go
|
package tracing
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/grafana/dskit/modules"
|
||
|
|
"github.com/grafana/dskit/services"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestModuleManagerWrapper_RegisterModule(t *testing.T) {
|
||
|
|
t.Run("registers module with wrapped init function", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
// Set context so getContext doesn't block
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
called := false
|
||
|
|
mockService := &mockNamedService{name: "test-service"}
|
||
|
|
initFn := func() (services.Service, error) {
|
||
|
|
called = true
|
||
|
|
return mockService, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
wrapper.RegisterModule("test-module", initFn)
|
||
|
|
|
||
|
|
// Verify the module was registered
|
||
|
|
require.True(t, manager.IsModuleRegistered("test-module"))
|
||
|
|
|
||
|
|
// Initialize the module to test the wrapped init function
|
||
|
|
service, err := manager.InitModuleServices("test-module")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.True(t, called)
|
||
|
|
require.NotNil(t, service)
|
||
|
|
|
||
|
|
// Verify listener was added to the service
|
||
|
|
require.Len(t, mockService.listeners, 1)
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("propagates init function errors", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
// Set context so getContext doesn't block
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
expectedErr := errors.New("init error")
|
||
|
|
initFn := func() (services.Service, error) {
|
||
|
|
return nil, expectedErr
|
||
|
|
}
|
||
|
|
|
||
|
|
wrapper.RegisterModule("test-module", initFn)
|
||
|
|
|
||
|
|
// Try to initialize the module
|
||
|
|
_, err := manager.InitModuleServices("test-module")
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Contains(t, err.Error(), expectedErr.Error())
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("handles nil init function", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
// Set context so getContext doesn't block
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
// Register module with nil init function
|
||
|
|
wrapper.RegisterModule("nil-module", nil)
|
||
|
|
|
||
|
|
// Verify the module was registered
|
||
|
|
require.True(t, manager.IsModuleRegistered("nil-module"))
|
||
|
|
|
||
|
|
// Initialize the module - should work with nil function
|
||
|
|
service, err := manager.InitModuleServices("nil-module")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Empty(t, service) // Should return empty map for nil function
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestModuleManagerWrapper_RegisterInvisibleModule(t *testing.T) {
|
||
|
|
t.Run("registers invisible module with wrapped init function", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
// Set context so getContext doesn't block
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
called := false
|
||
|
|
mockService := &mockNamedService{name: "invisible-service"}
|
||
|
|
initFn := func() (services.Service, error) {
|
||
|
|
called = true
|
||
|
|
return mockService, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
wrapper.RegisterInvisibleModule("invisible-module", initFn)
|
||
|
|
|
||
|
|
// Verify the module was registered
|
||
|
|
require.True(t, manager.IsModuleRegistered("invisible-module"))
|
||
|
|
|
||
|
|
// Initialize the module to test the wrapped init function
|
||
|
|
service, err := manager.InitModuleServices("invisible-module")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.True(t, called)
|
||
|
|
require.NotNil(t, service)
|
||
|
|
|
||
|
|
// Verify listener was added to the service
|
||
|
|
require.Len(t, mockService.listeners, 1)
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("handles nil init function", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
// Set context so getContext doesn't block
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
// Register invisible module with nil init function
|
||
|
|
wrapper.RegisterInvisibleModule("nil-invisible-module", nil)
|
||
|
|
|
||
|
|
// Verify the module was registered
|
||
|
|
require.True(t, manager.IsModuleRegistered("nil-invisible-module"))
|
||
|
|
|
||
|
|
// Initialize the module - should work with nil function
|
||
|
|
service, err := manager.InitModuleServices("nil-invisible-module")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Empty(t, service) // Should return empty map for nil function
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestModuleManagerWrapper_SetContext(t *testing.T) {
|
||
|
|
t.Run("sets context and closes ready channel", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
ctx := context.Background()
|
||
|
|
|
||
|
|
// Verify ready channel is not closed initially
|
||
|
|
select {
|
||
|
|
case <-wrapper.ready:
|
||
|
|
t.Fatal("ready channel should not be closed initially")
|
||
|
|
default:
|
||
|
|
// Expected
|
||
|
|
}
|
||
|
|
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
// Verify context is set and ready channel is closed
|
||
|
|
require.Equal(t, ctx, wrapper.ctx)
|
||
|
|
select {
|
||
|
|
case <-wrapper.ready:
|
||
|
|
// Expected - channel should be closed
|
||
|
|
case <-time.After(100 * time.Millisecond):
|
||
|
|
t.Fatal("ready channel should be closed after SetContext")
|
||
|
|
}
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("ignores subsequent calls", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
ctx1 := context.Background()
|
||
|
|
type contextKey string
|
||
|
|
ctx2 := context.WithValue(context.Background(), contextKey("key"), "value")
|
||
|
|
|
||
|
|
wrapper.SetContext(ctx1)
|
||
|
|
wrapper.SetContext(ctx2) // Should be ignored
|
||
|
|
|
||
|
|
require.Equal(t, ctx1, wrapper.ctx)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestModuleManagerWrapper_getContext(t *testing.T) {
|
||
|
|
t.Run("waits for ready channel and returns context", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
ctx := context.Background()
|
||
|
|
|
||
|
|
// Start goroutine to get context
|
||
|
|
resultCh := make(chan context.Context)
|
||
|
|
go func() {
|
||
|
|
resultCh <- wrapper.getContext()
|
||
|
|
}()
|
||
|
|
|
||
|
|
// Verify it's waiting
|
||
|
|
select {
|
||
|
|
case <-resultCh:
|
||
|
|
t.Fatal("getContext should wait for ready channel")
|
||
|
|
case <-time.After(50 * time.Millisecond):
|
||
|
|
// Expected - should be waiting
|
||
|
|
}
|
||
|
|
|
||
|
|
// Set context
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
// Verify getContext returns the context
|
||
|
|
select {
|
||
|
|
case result := <-resultCh:
|
||
|
|
require.Equal(t, ctx, result)
|
||
|
|
case <-time.After(100 * time.Millisecond):
|
||
|
|
t.Fatal("getContext should return after SetContext")
|
||
|
|
}
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestModuleManagerWrapper_wrapInitFn(t *testing.T) {
|
||
|
|
t.Run("adds listener to NamedService", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
mockService := &mockNamedService{name: "test-service"}
|
||
|
|
initFn := func() (services.Service, error) {
|
||
|
|
return mockService, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
wrappedFn := wrapper.wrapInitFn(initFn)
|
||
|
|
service, err := wrappedFn()
|
||
|
|
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, mockService, service)
|
||
|
|
|
||
|
|
// Verify listener was added to the service
|
||
|
|
require.Len(t, mockService.listeners, 1)
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("handles regular service without NamedService interface", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
mockService := &mockService{}
|
||
|
|
initFn := func() (services.Service, error) {
|
||
|
|
return mockService, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
wrappedFn := wrapper.wrapInitFn(initFn)
|
||
|
|
service, err := wrappedFn()
|
||
|
|
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Equal(t, mockService, service)
|
||
|
|
// No listener should be added for non-NamedService
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("propagates init function errors", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
expectedErr := errors.New("init error")
|
||
|
|
initFn := func() (services.Service, error) {
|
||
|
|
return nil, expectedErr
|
||
|
|
}
|
||
|
|
|
||
|
|
wrappedFn := wrapper.wrapInitFn(initFn)
|
||
|
|
service, err := wrappedFn()
|
||
|
|
|
||
|
|
require.Error(t, err)
|
||
|
|
require.Equal(t, expectedErr, err)
|
||
|
|
require.Nil(t, service)
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("handles nil service return", func(t *testing.T) {
|
||
|
|
manager := modules.NewManager(nil)
|
||
|
|
wrapper := WrapModuleManager(manager)
|
||
|
|
|
||
|
|
ctx := context.Background()
|
||
|
|
wrapper.SetContext(ctx)
|
||
|
|
|
||
|
|
initFn := func() (services.Service, error) {
|
||
|
|
return nil, nil // Return nil service with no error
|
||
|
|
}
|
||
|
|
|
||
|
|
wrappedFn := wrapper.wrapInitFn(initFn)
|
||
|
|
service, err := wrappedFn()
|
||
|
|
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.Nil(t, service)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
// Mock implementations for testing
|
||
|
|
|
||
|
|
type mockService struct {
|
||
|
|
listeners []services.Listener
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockService) AddListener(listener services.Listener) func() {
|
||
|
|
m.listeners = append(m.listeners, listener)
|
||
|
|
return func() {}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockService) AwaitRunning(ctx context.Context) error {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockService) AwaitTerminated(ctx context.Context) error {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockService) FailureCase() error {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockService) ServiceName() string {
|
||
|
|
return "mock-service"
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockService) StartAsync(ctx context.Context) error {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockService) State() services.State {
|
||
|
|
return services.New
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockService) StopAsync() {
|
||
|
|
}
|
||
|
|
|
||
|
|
type mockNamedService struct {
|
||
|
|
mockService
|
||
|
|
name string
|
||
|
|
}
|
||
|
|
|
||
|
|
func (m *mockNamedService) ServiceName() string {
|
||
|
|
return m.name
|
||
|
|
}
|