From 71e740b6e8b756f2419809068948d1832c3103ca Mon Sep 17 00:00:00 2001
From: Gorbunov Kirill Andreevich <kgorbunov@mtsbank.ru>
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)
+		})
+	}
+}