fix save return code in bin mode after writing result

server/mode_bin.go

@@ -81,9 +81,6 @@ func (bh *binHandler) ServeHTTP(writer http.ResponseWriter, request *http.Reques
 		codeReset.Status(http.StatusBadGateway)
 	}
 
-	if marker.dataSent {
-		return
-	}
 	if err != nil {
 		writer.Header().Set("X-Return-Code", strconv.Itoa(cmd.ProcessState.ExitCode()))
 		writer.WriteHeader(http.StatusBadGateway)

server/server_test.go

@@ -62,7 +62,7 @@ func testServer(t *testing.T, cfg server.Config, units map[string]server.Unit) *
 func Test_create(t *testing.T) {
 	srv := testServer(t, server.DefaultConfig(), map[string]server.Unit{
 		"hello": {
-			Command: "echo hello world",
+			Command: "echo -n hello world",
 		},
 	})
 	defer srv.Close()
@@ -96,7 +96,7 @@ func Test_create(t *testing.T) {
 			resultLocation = res.Header().Get("Location")
 			break
 		}
-		if assert.Equal(t, http.StatusNotFound, res.Code) {
+		if !assert.Equal(t, http.StatusNotFound, res.Code) {
 			return
 		}
 		time.Sleep(time.Second)
@@ -108,3 +108,42 @@ func Test_create(t *testing.T) {
 	assert.Equal(t, http.StatusOK, res.Code)
 	assert.Equal(t, "hello world", res.Body.String())
 }
+
+func Test_retryIfDataReturnedInBinMode(t *testing.T) {
+	srv := testServer(t, server.DefaultConfig(), map[string]server.Unit{
+		"hello": {
+			Command: "echo hello world; exit 1",
+		},
+	})
+	defer srv.Close()
+
+	req := httptest.NewRequest(http.MethodPost, "/hello/", bytes.NewBufferString("hello world"))
+	res := httptest.NewRecorder()
+	srv.ServeHTTP(res, req)
+	assert.Equal(t, http.StatusSeeOther, res.Code)
+	assert.NotEmpty(t, res.Header().Get("X-Correlation-Id"))
+	assert.Equal(t, "/hello/"+res.Header().Get("X-Correlation-Id"), res.Header().Get("Location"))
+	location := res.Header().Get("Location")
+
+	// wait for first result
+	for {
+		req = httptest.NewRequest(http.MethodGet, location, nil)
+		res = httptest.NewRecorder()
+		srv.ServeHTTP(res, req)
+		if !assert.Equal(t, http.StatusOK, res.Code) {
+			return
+		}
+		var info meta.Request
+		err := json.Unmarshal(res.Body.Bytes(), &info)
+		assert.NoError(t, err)
+		if len(info.Attempts) == 0 {
+			time.Sleep(time.Second)
+			continue
+		}
+		atp := info.Attempts[0]
+		assert.Equal(t, http.StatusBadGateway, atp.Code)
+		assert.Equal(t, "1", atp.Headers.Get("X-Return-Code"))
+		break
+	}
+
+}