Fix PDF download error (#115)

checkout · May 29, 2024 · e5ef57f · e5ef57f
1 parent dd6ea77
commit e5ef57f
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 4 deletions.
diff --git a/common/common.go b/common/common.go
@@ -142,9 +142,11 @@ type (
 		HttpMetadata HttpMetadata
 	}
 
+	Data interface{}
+
 	ContentResponse struct {
 		HttpMetadata HttpMetadata
-		Content      string `json:"content,omitempty"`
+		Content      Data `json:"content,omitempty"`
 	}
 
 	HttpMetadata struct {

diff --git a/common/serializer.go b/common/serializer.go
@@ -30,14 +30,15 @@ func Unmarshal(metadata *HttpMetadata, responseMapping interface{}) error {
 
 	if isContentTypeText(metadata) {
 		addContent(metadata.ResponseBody, responseMapping)
+	} else if isContentTypePDF(metadata) {
+		addPdfContent(metadata.ResponseBody, responseMapping)
 	} else {
 		if err := json.Unmarshal(metadata.ResponseBody, &responseMapping); err != nil {
 			return err
 		}
 	}
 
 	addHttpMetadata(metadata, responseMapping)
-
 	return nil
 }
 
@@ -48,13 +49,24 @@ func addHttpMetadata(metadata *HttpMetadata, response interface{}) {
 	}
 }
 
+func isContentTypeText(metadata *HttpMetadata) bool {
+	return strings.HasPrefix(metadata.Headers.Header.Get("Content-Type"), "text/csv")
+}
+
 func addContent(content []byte, response interface{}) {
 	v := reflect.ValueOf(response).Elem().FieldByName("Content")
 	if v.IsValid() {
 		v.Set(reflect.ValueOf(string(content)))
 	}
 }
 
-func isContentTypeText(metadata *HttpMetadata) bool {
-	return strings.HasPrefix(metadata.Headers.Header.Get("Content-Type"), "text/csv")
+func isContentTypePDF(metadata *HttpMetadata) bool {
+	return strings.HasPrefix(metadata.Headers.Header.Get("Content-Type"), "application/pdf")
+}
+
+func addPdfContent(content []byte, response interface{}) {
+	v := reflect.ValueOf(response).Elem().FieldByName("Content")
+	if v.IsValid() {
+		v.Set(reflect.ValueOf(content))
+	}
 }
diff --git a/common/serializer_test.go b/common/serializer_test.go
@@ -43,3 +43,53 @@ func TestMarshal(t *testing.T) {
 		})
 	}
 }
+
+func TestUnmarshalJSON(t *testing.T) {
+	responseBody := []byte(`{"type":"example","control_type":"example_control"}`)
+	headers := &Headers{Header: map[string][]string{"Content-Type": {"application/json"}}}
+	metadata := &HttpMetadata{ResponseBody: responseBody, Headers: headers}
+
+	var response TypeMapping
+	err := Unmarshal(metadata, &response)
+	if err != nil {
+		t.Fatalf("Unmarshal failed: %v", err)
+	}
+
+	assert.Equal(t, response, TypeMapping{Type: "example", ControlType: "example_control"})
+}
+
+func TestUnmarshalText(t *testing.T) {
+	responseBody := []byte("name,age\nJohn,30")
+	headers := &Headers{Header: map[string][]string{"Content-Type": {"text/csv"}}}
+	metadata := &HttpMetadata{ResponseBody: responseBody, Headers: headers}
+
+	type Response struct {
+		Content string
+	}
+
+	var response Response
+	err := Unmarshal(metadata, &response)
+	if err != nil {
+		t.Fatalf("Unmarshal failed: %v", err)
+	}
+
+	assert.Equal(t, response.Content, "name,age\nJohn,30")
+}
+
+func TestUnmarshalPDF(t *testing.T) {
+	responseBody := []byte("%PDF-1.4\n%...")
+	headers := &Headers{Header: map[string][]string{"Content-Type": {"application/pdf"}}}
+	metadata := &HttpMetadata{ResponseBody: responseBody, Headers: headers}
+
+	type Response struct {
+		Content []byte
+	}
+
+	var response Response
+	err := Unmarshal(metadata, &response)
+	if err != nil {
+		t.Fatalf("Unmarshal failed: %v", err)
+	}
+
+	assert.Equal(t, response.Content, responseBody)
+}