From d06ec675b8215326b41b41db1e47c30e3a81f3a9 Mon Sep 17 00:00:00 2001
From: Jordan Barrett <jordan.barrett@canonical.com>
Date: Thu, 3 Aug 2023 11:17:57 +0800
Subject: [PATCH] add test for encoding loki request

---
 internals/overlord/logstate/loki_test.go | 122 ++++++++++++++++++++++-
 1 file changed, 121 insertions(+), 1 deletion(-)

diff --git a/internals/overlord/logstate/loki_test.go b/internals/overlord/logstate/loki_test.go
index b393ef75d..2b843da2d 100644
--- a/internals/overlord/logstate/loki_test.go
+++ b/internals/overlord/logstate/loki_test.go
@@ -2,19 +2,127 @@ package logstate
 
 import (
 	"context"
+	"encoding/json"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"time"
 
+	. "gopkg.in/check.v1"
+
 	"github.com/canonical/pebble/internals/plan"
 	"github.com/canonical/pebble/internals/servicelog"
-	. "gopkg.in/check.v1"
 )
 
 type lokiSuite struct{}
 
 var _ = Suite(&lokiSuite{})
 
+func (*lokiSuite) TestLokiRequest(c *C) {
+	input := []servicelog.Entry{{
+		Time:    time.Date(2023, 12, 31, 12, 34, 50, 0, time.UTC),
+		Service: "svc1",
+		Message: "log line #1",
+	}, {
+		Time:    time.Date(2023, 12, 31, 12, 34, 51, 0, time.UTC),
+		Service: "svc2",
+		Message: "log line #2",
+	}, {
+		Time:    time.Date(2023, 12, 31, 12, 34, 52, 0, time.UTC),
+		Service: "svc1",
+		Message: "log line #3",
+	}, {
+		Time:    time.Date(2023, 12, 31, 12, 34, 53, 0, time.UTC),
+		Service: "svc3",
+		Message: "log line #4",
+	}, {
+		Time:    time.Date(2023, 12, 31, 12, 34, 54, 0, time.UTC),
+		Service: "svc1",
+		Message: "log line #5",
+	}, {
+		Time:    time.Date(2023, 12, 31, 12, 34, 55, 0, time.UTC),
+		Service: "svc3",
+		Message: "log line #6",
+	}, {
+		Time:    time.Date(2023, 12, 31, 12, 34, 56, 0, time.UTC),
+		Service: "svc2",
+		Message: "log line #7",
+	}, {
+		Time:    time.Date(2023, 12, 31, 12, 34, 57, 0, time.UTC),
+		Service: "svc1",
+		Message: "log line #8",
+	}, {
+		Time:    time.Date(2023, 12, 31, 12, 34, 58, 0, time.UTC),
+		Service: "svc4",
+		Message: "log line #9",
+	}}
+
+	expected := `
+{
+  "streams": [
+    {
+      "stream": {
+        "pebble_service": "svc1"
+      },
+      "values": [
+          [ "1704026090000000000", "log line #1" ],
+          [ "1704026092000000000", "log line #3" ],
+          [ "1704026094000000000", "log line #5" ],
+          [ "1704026097000000000", "log line #8" ]
+      ]
+    },
+    {
+      "stream": {
+        "pebble_service": "svc2"
+      },
+      "values": [
+          [ "1704026091000000000", "log line #2" ],
+          [ "1704026096000000000", "log line #7" ]
+      ]
+    },
+    {
+      "stream": {
+        "pebble_service": "svc3"
+      },
+      "values": [
+          [ "1704026093000000000", "log line #4" ],
+          [ "1704026095000000000", "log line #6" ]
+      ]
+    },
+    {
+      "stream": {
+        "pebble_service": "svc4"
+      },
+      "values": [
+          [ "1704026098000000000", "log line #9" ]
+      ]
+    }
+  ]
+}
+`
+
+	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		c.Assert(r.Method, Equals, http.MethodPost)
+		c.Assert(r.Header.Get("Content-Type"), Equals, "application/json")
+
+		reqBody, err := io.ReadAll(r.Body)
+		c.Assert(err, IsNil)
+		expFlattened, err := flattenJSON(expected)
+		c.Assert(err, IsNil)
+		c.Assert(string(reqBody), Equals, expFlattened)
+	}))
+	defer s.Close()
+
+	cl := newLokiClient(&plan.LogTarget{Location: s.URL})
+	for _, entry := range input {
+		err := cl.Write(context.Background(), entry)
+		c.Assert(err, IsNil)
+	}
+
+	err := cl.Flush(context.Background())
+	c.Assert(err, IsNil)
+}
+
 func (*lokiSuite) TestLokiFlushCancelContext(c *C) {
 	serverCtx, killServer := context.WithCancel(context.Background())
 	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -53,3 +161,15 @@ func (*lokiSuite) TestLokiFlushCancelContext(c *C) {
 		c.Fatal("lokiClient.Flush took too long to return after context timeout")
 	}
 }
+
+// Strips all extraneous whitespace from JSON
+func flattenJSON(s string) (string, error) {
+	var v any
+	err := json.Unmarshal([]byte(s), &v)
+	if err != nil {
+		return "", err
+	}
+
+	b, err := json.Marshal(v)
+	return string(b), err
+}