From 6aa5f967a5a7d6d49d1db893a90c00514b3e6302 Mon Sep 17 00:00:00 2001
From: Jordan Barrett <jordan.barrett@canonical.com>
Date: Thu, 3 Aug 2023 10:35:53 +0800
Subject: [PATCH] add TestLokiFlushCancelContext

---
 internals/overlord/logstate/loki_test.go | 55 ++++++++++++++++++++++++
 1 file changed, 55 insertions(+)
 create mode 100644 internals/overlord/logstate/loki_test.go

diff --git a/internals/overlord/logstate/loki_test.go b/internals/overlord/logstate/loki_test.go
new file mode 100644
index 000000000..b393ef75d
--- /dev/null
+++ b/internals/overlord/logstate/loki_test.go
@@ -0,0 +1,55 @@
+package logstate
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"time"
+
+	"github.com/canonical/pebble/internals/plan"
+	"github.com/canonical/pebble/internals/servicelog"
+	. "gopkg.in/check.v1"
+)
+
+type lokiSuite struct{}
+
+var _ = Suite(&lokiSuite{})
+
+func (*lokiSuite) TestLokiFlushCancelContext(c *C) {
+	serverCtx, killServer := context.WithCancel(context.Background())
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		select {
+		case <-serverCtx.Done():
+		// Simulate a slow-responding server
+		case <-time.After(10 * time.Second):
+		}
+	}))
+	defer s.Close()
+	defer killServer()
+
+	cl := newLokiClient(&plan.LogTarget{Location: s.URL})
+	err := cl.Write(context.Background(), servicelog.Entry{
+		Time:    time.Now(),
+		Service: "svc1",
+		Message: "this is a log line\n",
+	})
+	c.Assert(err, IsNil)
+
+	// Cancel the Flush context quickly
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Millisecond)
+	defer cancel()
+
+	flushReturned := make(chan struct{})
+	go func() {
+		err = cl.Flush(ctx)
+		close(flushReturned)
+		c.Assert(err, ErrorMatches, ".*context deadline exceeded")
+	}()
+
+	// Check Flush returns quickly after context timeout
+	select {
+	case <-flushReturned:
+	case <-time.After(2 * time.Millisecond):
+		c.Fatal("lokiClient.Flush took too long to return after context timeout")
+	}
+}