From 56da86258605983bc670f1ec98863ae1e520286e Mon Sep 17 00:00:00 2001
From: Abi Hafshin <abi@hafs.in>
Date: Mon, 6 Dec 2021 20:49:21 +0700
Subject: [PATCH] add writer test

---
 writer_test.go | 77 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)
 create mode 100644 writer_test.go

diff --git a/writer_test.go b/writer_test.go
new file mode 100644
index 0000000..0e7ab66
--- /dev/null
+++ b/writer_test.go
@@ -0,0 +1,77 @@
+package compress
+
+import (
+	"bytes"
+	"context"
+	"net/http"
+	"testing"
+
+	"gotest.tools/assert"
+)
+
+func Test_withKnowContentEncoding(t *testing.T) {
+	w := newMockRW()
+	c := newConfig()
+	encoding := "gzip"
+	rw := responseWriter{
+		ResponseWriter: w,
+		c:              c,
+		ctx:            context.Background(),
+		factory:        c.encoders[encoding].factory,
+		encoding:       encoding,
+	}
+
+	rw.Header().Set("content-encoding", "flate")
+	rw.Header().Set("content-type", "text/plain")
+	rw.WriteHeader(200)
+
+	assert.Equal(t, rw.dontEncode, true)
+	assert.Equal(t, rw.shouldEncode, false)
+}
+
+func Test_withKnowContentLength(t *testing.T) {
+	w := newMockRW()
+	c := newConfig()
+	encoding := "gzip"
+	rw := responseWriter{
+		ResponseWriter: w,
+		c:              c,
+		ctx:            context.Background(),
+		factory:        c.encoders[encoding].factory,
+		encoding:       encoding,
+	}
+
+	rw.Header().Set("content-type", "text/plain")
+	rw.Header().Set("content-length", "100000")
+	rw.WriteHeader(400)
+
+	assert.Equal(t, rw.dontEncode, false)
+	assert.Equal(t, rw.shouldEncode, true)
+	assert.Equal(t, rw.status, 400)
+}
+
+type mockResponseWriter struct {
+	h      http.Header
+	status int
+	body   bytes.Buffer
+}
+
+func newMockRW() http.ResponseWriter {
+	return &mockResponseWriter{
+		status: 200,
+		body:   bytes.Buffer{},
+		h:      http.Header{},
+	}
+}
+
+func (w *mockResponseWriter) Header() http.Header {
+	return w.h
+}
+
+func (w *mockResponseWriter) Write(c []byte) (int, error) {
+	return w.body.Write(c)
+}
+
+func (w *mockResponseWriter) WriteHeader(status int) {
+	w.status = status
+}