From 4b95f611c2c7a214fc9b9f2415e6c0851ac58476 Mon Sep 17 00:00:00 2001
From: Todd Treece <360020+toddtreece@users.noreply.github.com>
Date: Mon, 24 Jul 2023 14:01:07 -0400
Subject: [PATCH] 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           |  9 +++++++++
 pkg/modules/util.go              | 14 +++++++++++---
 pkg/server/server.go             |  5 +++++
 pkg/tests/testinfra/testinfra.go |  2 ++
 4 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/pkg/modules/modules.go b/pkg/modules/modules.go
index d05c08708bc..26321f156dc 100644
--- a/pkg/modules/modules.go
+++ b/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
diff --git a/pkg/modules/util.go b/pkg/modules/util.go
index 39a2449a5b6..f2b12e0561d 100644
--- a/pkg/modules/util.go
+++ b/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 {
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 531d2467b5d..29a238b1029 100644
--- a/pkg/server/server.go
+++ b/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 {
diff --git a/pkg/tests/testinfra/testinfra.go b/pkg/tests/testinfra/testinfra.go
index baf1eab1c08..a0fbd2ead22 100644
--- a/pkg/tests/testinfra/testinfra.go
+++ b/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)