Add download buttons for spam and ham samples

umputun · Jan 21, 2024 · 89dd4a0 · 89dd4a0
1 parent 686c4b7
commit 89dd4a0
Show file tree

Hide file tree

Showing 4 changed files with 76 additions and 19 deletions.
diff --git a/app/webapi/assets/components/samples_list.html b/app/webapi/assets/components/samples_list.html
@@ -1,6 +1,9 @@
 <div class="row" id="samples-list">
     <div class="col-md-6">
-        <h4>Spam Samples ({{.TotalSpamSamples}})</h4>
+        <h4>
+            Spam Samples ({{.TotalSpamSamples}})
+            <a href="/download/spam" class="btn btn-custom-blue btn-sm">Download</a>
+        </h4>
         <ul class="list-group" id="spam-samples-list">
             {{range .SpamSamples}}
             <li class="list-group-item d-flex justify-content-between align-items-center">
@@ -19,7 +22,10 @@ <h4>Spam Samples ({{.TotalSpamSamples}})</h4>
     </div>
 
     <div class="col-md-6">
-        <h4>Ham Samples ({{.TotalHamSamples}})</h4>
+        <h4>
+            Ham Samples ({{.TotalSpamSamples}})
+            <a href="/download/ham" class="btn btn-custom-blue btn-sm">Download</a>
+        </h4>
         <ul class="list-group" id="ham-samples-list">
             {{range .HamSamples}}
             <li class="list-group-item d-flex justify-content-between align-items-center">

diff --git a/app/webapi/assets/styles.css b/app/webapi/assets/styles.css
@@ -30,3 +30,15 @@ body, #result{
     max-width: 500px;
     overflow: hidden;
 }
+
+.btn-custom-blue {
+    color: #eee;
+    background-color: #286ea7;
+    border-color: #3d86e5;
+}
+
+.btn-custom-blue:hover {
+    color: #fff;
+    background-color: #3286cb;
+    border-color: #3d86e5;
+}
diff --git a/app/webapi/webapi.go b/app/webapi/webapi.go
@@ -138,6 +138,15 @@ func (s *Server) routes(router *chi.Mux) *chi.Mux {
 			r.Post("/ham", s.deleteSampleHandler(s.SpamFilter.RemoveDynamicHamSample))
 		})
 
+		authApi.Route("/download", func(r chi.Router) {
+			r.Get("/spam", s.downloadSampleHandler(func(spam, ham []string) ([]string, string) {
+				return spam, "spam.txt"
+			}))
+			r.Get("/ham", s.downloadSampleHandler(func(spam, ham []string) ([]string, string) {
+				return ham, "ham.txt"
+			}))
+		})
+
 		authApi.Get("/samples", s.getDynamicSamplesHandler)    // get dynamic samples
 		authApi.Put("/samples", s.reloadDynamicSamplesHandler) // reload samples
 
@@ -237,6 +246,25 @@ func (s *Server) getDynamicSamplesHandler(w http.ResponseWriter, _ *http.Request
 	rest.RenderJSON(w, rest.JSON{"spam": spam, "ham": ham})
 }
 
+// downloadSampleHandler handles GET /download/spam|ham request. It returns dynamic samples both for spam and ham.
+func (s *Server) downloadSampleHandler(pickFn func(spam, ham []string) ([]string, string)) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		spam, ham, err := s.SpamFilter.DynamicSamples()
+		if err != nil {
+			w.WriteHeader(http.StatusInternalServerError)
+			rest.RenderJSON(w, rest.JSON{"error": "can't get dynamic samples", "details": err.Error()})
+			return
+		}
+		samples, name := pickFn(spam, ham)
+		body := strings.Join(samples, "\n")
+		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
+		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", name))
+		w.Header().Set("Content-Length", strconv.Itoa(len(body)))
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(body))
+	}
+}
+
 // updateSampleHandler handles POST /update/spam|ham request. It updates dynamic samples both for spam and ham.
 func (s *Server) updateSampleHandler(updFn func(msg string) error) func(w http.ResponseWriter, r *http.Request) {
 	return func(w http.ResponseWriter, r *http.Request) {

diff --git a/app/webapi/webapi_test.go b/app/webapi/webapi_test.go
@@ -540,8 +540,8 @@ func TestServer_deleteSampleHandler(t *testing.T) {
 		assert.Equal(t, http.StatusOK, rr.Code, "handler returned wrong status code")
 		body := rr.Body.String()
 		t.Log(body)
-		assert.Contains(t, body, "<h4>Spam Samples (2)</h4>", "response should contain spam samples")
-		assert.Contains(t, body, "<h4>Ham Samples (2)</h4>", "response should contain ham samples")
+		assert.Contains(t, body, "Spam Samples (2)", "response should contain spam samples")
+		assert.Contains(t, body, "Ham Samples (2)", "response should contain ham samples")
 		require.Equal(t, 1, len(spamFilterMock.RemoveDynamicHamSampleCalls()))
 		assert.Equal(t, "test message", spamFilterMock.RemoveDynamicHamSampleCalls()[0].Sample)
 	})
@@ -878,7 +878,7 @@ func TestServer_logoHandler(t *testing.T) {
 	assert.Equal(t, "image/png", rr.Header().Get("Content-Type"), "handler should return CSS content type")
 }
 
-func TestServer_getDynamicSamplesHandler(t *testing.T) {
+func Test_downloadSampleHandler(t *testing.T) {
 	mockSpamFilter := &mocks.SpamFilterMock{
 		DynamicSamplesFunc: func() ([]string, []string, error) {
 			return []string{"spam1", "spam2"}, []string{"ham1", "ham2"}, nil
@@ -889,37 +889,48 @@ func TestServer_getDynamicSamplesHandler(t *testing.T) {
 		SpamFilter: mockSpamFilter,
 	})
 
-	t.Run("successful response", func(t *testing.T) {
-		req, err := http.NewRequest("GET", "/samples", http.NoBody)
+	t.Run("successful spam response", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "/download/spam", http.NoBody)
 		require.NoError(t, err)
 
 		rr := httptest.NewRecorder()
-		handler := http.HandlerFunc(server.getDynamicSamplesHandler)
+		handler := http.HandlerFunc(server.downloadSampleHandler(func(spam, ham []string) ([]string, string) {
+			return spam, "spam.txt"
+		}))
 
 		handler.ServeHTTP(rr, req)
 		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, "text/plain; charset=utf-8", rr.Header().Get("Content-Type"))
+		assert.Contains(t, rr.Header().Get("Content-Disposition"), "attachment; filename=\"spam.txt\"")
+	})
 
-		var response struct {
-			Spam []string `json:"spam"`
-			Ham  []string `json:"ham"`
-		}
-		err = json.Unmarshal(rr.Body.Bytes(), &response)
+	t.Run("successful ham response", func(t *testing.T) {
+		req, err := http.NewRequest("GET", "/download/ham", http.NoBody)
 		require.NoError(t, err)
 
-		assert.Equal(t, []string{"spam1", "spam2"}, response.Spam)
-		assert.Equal(t, []string{"ham1", "ham2"}, response.Ham)
+		rr := httptest.NewRecorder()
+		handler := http.HandlerFunc(server.downloadSampleHandler(func(spam, ham []string) ([]string, string) {
+			return spam, "ham.txt"
+		}))
+
+		handler.ServeHTTP(rr, req)
+		assert.Equal(t, http.StatusOK, rr.Code)
+		assert.Equal(t, "text/plain; charset=utf-8", rr.Header().Get("Content-Type"))
+		assert.Contains(t, rr.Header().Get("Content-Disposition"), "attachment; filename=\"ham.txt\"")
 	})
 
 	t.Run("error handling", func(t *testing.T) {
 		mockSpamFilter.DynamicSamplesFunc = func() ([]string, []string, error) {
 			return nil, nil, errors.New("test error")
 		}
 
-		req, err := http.NewRequest("GET", "/samples", http.NoBody)
+		req, err := http.NewRequest("GET", "/download/ham", http.NoBody)
 		require.NoError(t, err)
 
 		rr := httptest.NewRecorder()
-		handler := http.HandlerFunc(server.getDynamicSamplesHandler)
+		handler := http.HandlerFunc(server.downloadSampleHandler(func(spam, ham []string) ([]string, string) {
+			return spam, "ham.txt"
+		}))
 
 		handler.ServeHTTP(rr, req)
 		assert.Equal(t, http.StatusInternalServerError, rr.Code)
@@ -1048,10 +1059,10 @@ func TestServer_renderSamples(t *testing.T) {
 	assert.Equal(t, http.StatusOK, w.Code)
 	assert.Equal(t, "text/html; charset=utf-8", w.Header().Get("Content-Type"))
 	t.Log(w.Body.String())
-	assert.Contains(t, w.Body.String(), "<h4>Spam Samples (2)</h4>")
+	assert.Contains(t, w.Body.String(), "Spam Samples (2)")
 	assert.Contains(t, w.Body.String(), "spam1")
 	assert.Contains(t, w.Body.String(), "spam2")
-	assert.Contains(t, w.Body.String(), "<h4>Ham Samples (2)</h4>")
+	assert.Contains(t, w.Body.String(), "Ham Samples (2)")
 	assert.Contains(t, w.Body.String(), "ham1")
 	assert.Contains(t, w.Body.String(), "ham2")
 }