From d30befa2d7d79cf07a6212e486c62b4a145268f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Fri, 15 Apr 2022 16:14:17 +0200
Subject: [PATCH] Manager: add assert function for testing JSON responses

This makes it much easier to test an API response actually matches the
expected JSON values.
---
 internal/manager/api_impl/support_test.go | 37 +++++++++++++++++++++++
 internal/manager/api_impl/workers_test.go | 31 +++++++++----------
 2 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/internal/manager/api_impl/support_test.go b/internal/manager/api_impl/support_test.go
index 93c96877..06134d8a 100644
--- a/internal/manager/api_impl/support_test.go
+++ b/internal/manager/api_impl/support_test.go
@@ -5,12 +5,15 @@ package api_impl
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"testing"
 
 	"github.com/golang/mock/gomock"
 	"github.com/labstack/echo/v4"
+	"github.com/stretchr/testify/assert"
 	"gorm.io/gorm"
 
 	"git.blender.org/flamenco/internal/manager/api_impl/mocks"
@@ -64,6 +67,7 @@ func (mf *mockedFlamenco) prepareMockedJSONRequest(requestBody interface{}) echo
 }
 
 // prepareMockedJSONRequest returns an `echo.Context` that has an empty request body attached to it.
+// `body` may be `nil` to indicate "no body".
 func (mf *mockedFlamenco) prepareMockedRequest(body io.Reader) echo.Context {
 	e := echo.New()
 
@@ -74,6 +78,39 @@ func (mf *mockedFlamenco) prepareMockedRequest(body io.Reader) echo.Context {
 	return c
 }
 
+func getRecordedResponse(echoCtx echo.Context) *http.Response {
+	writer := echoCtx.Response().Writer
+	resp, ok := writer.(*httptest.ResponseRecorder)
+	if !ok {
+		panic(fmt.Sprintf("response writer was not a `*httptest.ResponseRecorder` but a %T", writer))
+	}
+	return resp.Result()
+}
+
+// assertJSONResponse asserts that a recorded response is JSON with the given HTTP status code.
+func assertJSONResponse(t *testing.T, echoCtx echo.Context, expectStatusCode int, expectBody interface{}) {
+	resp := getRecordedResponse(echoCtx)
+	assert.Equal(t, expectStatusCode, resp.StatusCode)
+	contentType := resp.Header.Get(echo.HeaderContentType)
+
+	if !assert.Equal(t, "application/json; charset=UTF-8", contentType) {
+		t.Fatalf("response not JSON but %q, not going to compare body", contentType)
+		return
+	}
+
+	expectJSON, err := json.Marshal(expectBody)
+	if !assert.NoError(t, err) {
+		t.FailNow()
+	}
+
+	actualJSON, err := io.ReadAll(resp.Body)
+	if !assert.NoError(t, err) {
+		t.FailNow()
+	}
+
+	assert.JSONEq(t, string(expectJSON), string(actualJSON))
+}
+
 func testWorker() persistence.Worker {
 	return persistence.Worker{
 		Model:              gorm.Model{ID: 1},
diff --git a/internal/manager/api_impl/workers_test.go b/internal/manager/api_impl/workers_test.go
index 90d48540..d0b6cffd 100644
--- a/internal/manager/api_impl/workers_test.go
+++ b/internal/manager/api_impl/workers_test.go
@@ -4,9 +4,7 @@ package api_impl
 
 import (
 	"context"
-	"encoding/json"
 	"net/http"
-	"net/http/httptest"
 	"testing"
 
 	"github.com/golang/mock/gomock"
@@ -39,9 +37,15 @@ func TestTaskScheduleHappy(t *testing.T) {
 	err := mf.flamenco.ScheduleTask(echo)
 	assert.NoError(t, err)
 
-	resp := echo.Response().Writer.(*httptest.ResponseRecorder)
-	assert.Equal(t, http.StatusOK, resp.Code)
-	// TODO: check that the returned JSON actually matches what we expect.
+	// Check the response
+	assignedTask := api.AssignedTask{
+		Uuid:     task.UUID,
+		Job:      job.UUID,
+		Commands: []api.Command{},
+	}
+	assertJSONResponse(t, echo, http.StatusOK, assignedTask)
+	resp := getRecordedResponse(echo)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
 }
 
 func TestTaskScheduleNonActiveStatus(t *testing.T) {
@@ -60,8 +64,8 @@ func TestTaskScheduleNonActiveStatus(t *testing.T) {
 	err := mf.flamenco.ScheduleTask(echoCtx)
 	assert.NoError(t, err)
 
-	resp := echoCtx.Response().Writer.(*httptest.ResponseRecorder)
-	assert.Equal(t, http.StatusConflict, resp.Code)
+	resp := getRecordedResponse(echoCtx)
+	assert.Equal(t, http.StatusConflict, resp.StatusCode)
 }
 
 func TestTaskScheduleOtherStatusRequested(t *testing.T) {
@@ -80,13 +84,8 @@ func TestTaskScheduleOtherStatusRequested(t *testing.T) {
 	err := mf.flamenco.ScheduleTask(echoCtx)
 	assert.NoError(t, err)
 
-	resp := echoCtx.Response().Writer.(*httptest.ResponseRecorder)
-	assert.Equal(t, http.StatusLocked, resp.Code)
-
-	responseBody := api.WorkerStateChange{}
-	err = json.Unmarshal(resp.Body.Bytes(), &responseBody)
-	assert.NoError(t, err)
-	assert.Equal(t, worker.StatusRequested, responseBody.StatusRequested)
+	expectBody := api.WorkerStateChange{StatusRequested: api.WorkerStatusAsleep}
+	assertJSONResponse(t, echoCtx, http.StatusLocked, expectBody)
 }
 
 func TestWorkerSignoffTaskRequeue(t *testing.T) {
@@ -137,6 +136,6 @@ func TestWorkerSignoffTaskRequeue(t *testing.T) {
 	err := mf.flamenco.SignOff(echo)
 	assert.NoError(t, err)
 
-	resp := echo.Response().Writer.(*httptest.ResponseRecorder)
-	assert.Equal(t, http.StatusNoContent, resp.Code)
+	resp := getRecordedResponse(echo)
+	assert.Equal(t, http.StatusNoContent, resp.StatusCode)
 }