From 1329ea165ca357a481b9bf9a25352dd2caafcbef Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Fri, 8 Apr 2022 14:20:14 +0100
Subject: [PATCH] TestPromtail: avoid hanging when test fails midway (#5837)

* Allow Promtail.Shutdown to be called again

I want to do this in a test

* TestPromtail: avoid hanging when test fails midway

Also call Shutdown via defer so we clean up those resources when the test fails
---
 clients/pkg/promtail/promtail.go      | 3 +++
 clients/pkg/promtail/promtail_test.go | 4 ++++
 2 files changed, 7 insertions(+)

diff --git a/clients/pkg/promtail/promtail.go b/clients/pkg/promtail/promtail.go
index e7eca54dbe..3362a7414e 100644
--- a/clients/pkg/promtail/promtail.go
+++ b/clients/pkg/promtail/promtail.go
@@ -114,6 +114,9 @@ func (p *Promtail) Client() client.Client {
 func (p *Promtail) Shutdown() {
 	p.mtx.Lock()
 	defer p.mtx.Unlock()
+	if p.stopped {
+		return
+	}
 	p.stopped = true
 	if p.server != nil {
 		p.server.Shutdown()
diff --git a/clients/pkg/promtail/promtail_test.go b/clients/pkg/promtail/promtail_test.go
index 2f5eb746db..0b4852d22e 100644
--- a/clients/pkg/promtail/promtail_test.go
+++ b/clients/pkg/promtail/promtail_test.go
@@ -86,6 +86,9 @@ func TestPromtail(t *testing.T) {
 		server    = &http.Server{Addr: "localhost:3100", Handler: nil}
 	)
 	defer func() {
+		if t.Failed() {
+			return // Test has already failed; don't wait for everything to shut down.
+		}
 		fmt.Fprintf(os.Stdout, "wait close")
 		wg.Wait()
 		if err != nil {
@@ -117,6 +120,7 @@ func TestPromtail(t *testing.T) {
 			err = errors.Wrap(err, "Failed to start promtail")
 		}
 	}()
+	defer p.Shutdown() // In case the test fails before the call to Shutdown below.
 
 	expectedCounts := map[string]int{}