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") + } +}