From dd1f3d573e41d94653c1d1e9fbebdd177ce6c6ee Mon Sep 17 00:00:00 2001
From: Yamagishi Kazutoshi <ykzts@desire.sh>
Date: Fri, 21 May 2021 05:26:39 +0900
Subject: [PATCH] fix(transport): fix duplicate variables (#464)

---
 .github/workflows/go.yml |  3 +-
 transport.go             | 13 ++++++---
 transport_test.go        | 59 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
index a5c32fa2..baddba71 100644
--- a/.github/workflows/go.yml
+++ b/.github/workflows/go.yml
@@ -42,9 +42,10 @@ jobs:
       - name: Test
         env:
           GO111MODULE: on
+          MANAEL_ENABLE_AVIF: true
         run: |
           mkdir -p cover
-          go test -race -coverprofile=coverage.txt -covermode=atomic
+          go test -race -coverprofile=coverage.txt -covermode=atomic -v
 
       - uses: codecov/codecov-action@v1
         with:
diff --git a/transport.go b/transport.go
index f5a3b374..039da1e5 100644
--- a/transport.go
+++ b/transport.go
@@ -72,14 +72,19 @@ func (t *Transport) makeRequest(r *http.Request) (*http.Request, error) {
 	return r2, nil
 }
 
-func scanAcceptHeader(r *http.Request, t string) string {
-	f := os.Getenv("MANAEL_ENABLE_AVIF")
+func avifEnabled(w *http.Response) bool {
+	t := w.Header.Get("Content-Type")
+
+	return os.Getenv("MANAEL_ENABLE_AVIF") == "true" && t != "image/png"
+}
+
+func scanAcceptHeader(w *http.Response, r *http.Request) string {
 	a := r.Header.Get("Accept")
 
 	for _, v := range strings.Split(a, ",") {
 		t := strings.TrimSpace(v)
 
-		if f == "true" && t != "image/png" && strings.HasPrefix(t, "image/avif") {
+		if avifEnabled(w) && strings.HasPrefix(t, "image/avif") {
 			return "image/avif"
 		} else if strings.HasPrefix(t, "image/webp") {
 			return "image/webp"
@@ -112,7 +117,7 @@ func check(w *http.Response, r *http.Request) string {
 		return "*/*"
 	}
 
-	return scanAcceptHeader(r, t)
+	return scanAcceptHeader(w, r)
 }
 
 func convert(src io.Reader, t string) (*bytes.Buffer, error) {
diff --git a/transport_test.go b/transport_test.go
index 9e963031..70230d93 100644
--- a/transport_test.go
+++ b/transport_test.go
@@ -28,6 +28,7 @@ import (
 	"net"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"testing"
 	"time"
 
@@ -452,3 +453,61 @@ func TestTransport_RoundTrip_xForwardedFor(t *testing.T) {
 		}
 	}
 }
+
+var transportTests7 = []struct {
+	accept      string
+	path        string
+	statusCode  int
+	contentType string
+}{
+	{
+		"image/avif,image/webp,image/*,*/*;q=0.8",
+		"/photo.jpeg",
+		http.StatusOK,
+		"image/avif",
+	},
+	{
+		"image/avif,image/webp,image/*,*/*;q=0.8",
+		"/logo.png",
+		http.StatusOK,
+		"image/webp",
+	},
+}
+
+func TestTransport_RoundTrip_avif(t *testing.T) {
+	if os.Getenv("MANAEL_ENABLE_AVIF") != "true" {
+		t.Skip("Skipping test when avif disabled.")
+	}
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/logo.png", func(w http.ResponseWriter, r *http.Request) {
+		http.ServeFile(w, r, "testdata/logo.png")
+	})
+	mux.HandleFunc("/photo.jpeg", func(w http.ResponseWriter, r *http.Request) {
+		http.ServeFile(w, r, "testdata/photo.jpeg")
+	})
+
+	ts := httptest.NewServer(mux)
+	defer ts.Close()
+
+	tr := &manael.Transport{ts.URL, http.DefaultTransport}
+
+	for _, tc := range transportTests7 {
+		req := httptest.NewRequest(http.MethodGet, "https://manael.test"+tc.path, nil)
+		req.Header.Set("Accept", tc.accept)
+
+		resp, err := tr.RoundTrip(req)
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer resp.Body.Close()
+
+		if got, want := resp.StatusCode, tc.statusCode; got != want {
+			t.Errorf("Status Code is %d, want %d", got, want)
+		}
+
+		if got, want := resp.Header.Get("Content-Type"), tc.contentType; got != want {
+			t.Errorf("Content-Type is %s, want %s", got, want)
+		}
+	}
+}