Skip to content

Commit

Permalink
add test for encoding loki request
Browse files Browse the repository at this point in the history
  • Loading branch information
barrettj12 committed Aug 3, 2023
1 parent 6aa5f96 commit d06ec67
Showing 1 changed file with 121 additions and 1 deletion.
122 changes: 121 additions & 1 deletion internals/overlord/logstate/loki_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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
}

0 comments on commit d06ec67

Please sign in to comment.