Chore: Add AwaitHealthy to ModuleEngine and Server (#72215)

* Chore: Add AwaitHealthy to ModuleEngine and Server

* switch from fmt.Errorf to errors.New

Co-authored-by: Will Browne <wbrowne@users.noreply.github.com>

---------

Co-authored-by: Will Browne <wbrowne@users.noreply.github.com>

pkg/modules/modules.go

@@ -15,6 +15,7 @@ import (
 )
 
 type Engine interface {
+	AwaitHealthy(context.Context) error
 	Init(context.Context) error
 	Run(context.Context) error
 	Shutdown(context.Context) error
@@ -59,6 +60,14 @@ func ProvideService(
 	}
 }
 
+// AwaitHealthy waits for all registered modules to be healthy.
+func (m *service) AwaitHealthy(ctx context.Context) error {
+	if m.serviceManager == nil {
+		return errors.New("service manager has not been initialized")
+	}
+	return m.serviceManager.AwaitHealthy(ctx)
+}
+
 // Init initializes all registered modules.
 func (m *service) Init(_ context.Context) error {
 	var err error

pkg/modules/util.go

@@ -27,9 +27,17 @@ func (m *MockModuleManager) RegisterInvisibleModule(name string, initFn func() (
 }
 
 type MockModuleEngine struct {
-	InitFunc     func(context.Context) error
-	RunFunc      func(context.Context) error
-	ShutdownFunc func(context.Context) error
+	AwaitHealthyFunc func(context.Context) error
+	InitFunc         func(context.Context) error
+	RunFunc          func(context.Context) error
+	ShutdownFunc     func(context.Context) error
+}
+
+func (m *MockModuleEngine) AwaitHealthy(ctx context.Context) error {
+	if m.AwaitHealthyFunc != nil {
+		return m.AwaitHealthyFunc(ctx)
+	}
+	return nil
 }
 
 func (m *MockModuleEngine) Init(ctx context.Context) error {

pkg/server/server.go

@@ -112,6 +112,11 @@ func (s *Server) init(ctx context.Context) error {
 	return s.roleRegistry.RegisterFixedRoles(ctx)
 }
 
+// AwaitHealthy waits for the server to become healthy.
+func (s *Server) AwaitHealthy(ctx context.Context) error {
+	return s.moduleService.AwaitHealthy(ctx)
+}
+
 // Run initializes and starts services. This will block until all services have
 // exited. To initiate shutdown, call the Shutdown method in another goroutine.
 func (s *Server) Run(ctx context.Context) error {

pkg/tests/testinfra/testinfra.go

@@ -71,6 +71,8 @@ func StartGrafanaEnv(t *testing.T, grafDir, cfgPath string) (string, *server.Tes
 	})
 
 	// Wait for Grafana to be ready
+	err = env.Server.AwaitHealthy(ctx)
+	require.NoError(t, err)
 	addr := listener.Addr().String()
 	resp, err := http.Get(fmt.Sprintf("http://%s/api/health", addr))
 	require.NoError(t, err)