From 71e740b6e8b756f2419809068948d1832c3103ca Mon Sep 17 00:00:00 2001 From: Gorbunov Kirill Andreevich Date: Mon, 18 Mar 2024 16:00:49 +0300 Subject: [PATCH] add Flush --- handler/meter/meter.go | 2 ++ handler/meter/meter_test.go | 49 +++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) create mode 100644 handler/meter/meter_test.go diff --git a/handler/meter/meter.go b/handler/meter/meter.go index 4157cd9..3cbec16 100644 --- a/handler/meter/meter.go +++ b/handler/meter/meter.go @@ -92,6 +92,7 @@ func (h *Handler) Metrics(ctx context.Context, req *codecpb.Frame, rsp *codecpb. if md, ok := metadata.FromContext(ctx); gzipAccepted(md) && ok { md.Set(contentEncodingHeader, "gzip") + ctx = metadata.NewIncomingContext(ctx, md) gz := gzipPool.Get().(*gzip.Writer) defer gzipPool.Put(gz) @@ -99,6 +100,7 @@ func (h *Handler) Metrics(ctx context.Context, req *codecpb.Frame, rsp *codecpb. defer gz.Close() w = gz + gz.Flush() } if err := h.opts.Meter.Write(w, h.opts.MeterOptions...); err != nil { diff --git a/handler/meter/meter_test.go b/handler/meter/meter_test.go new file mode 100644 index 0000000..a3b7f51 --- /dev/null +++ b/handler/meter/meter_test.go @@ -0,0 +1,49 @@ +package meter + +import ( + "context" + "testing" + + codecpb "go.unistack.org/micro-proto/v4/codec" +) + +func TestHandler_Metrics(t *testing.T) { + type fields struct { + opts Options + } + type args struct { + ctx context.Context + req *codecpb.Frame + rsp *codecpb.Frame + } + tests := []struct { + name string + fields fields + args args + wantErr bool + }{ + { + "Test #1", + fields{ + opts: NewOptions(), + }, + args{ + context.Background(), + &codecpb.Frame{Data: []byte("gzip")}, + &codecpb.Frame{}, + }, + true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + h := &Handler{ + opts: tt.fields.opts, + } + if err := h.Metrics(tt.args.ctx, tt.args.req, tt.args.rsp); (err != nil) != tt.wantErr { + t.Errorf("Metrics() error = %v, wantErr %v", err, tt.wantErr) + } + t.Logf("RSP: %v", tt.args.rsp.Data) + }) + } +}