Update updater_test.go

digitalocean · Jun 12, 2024 · dc37bb8 · dc37bb8
1 parent ddd6ef7
commit dc37bb8
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/internal/metadata/updater/updater_test.go b/internal/metadata/updater/updater_test.go
@@ -24,12 +24,12 @@ func Test_agentInfoUpdaterImpl_Update(t *testing.T) {
 	}
 	tests := []struct {
 		name         string
-		expectations func(client *MockhttpClient)
+		expectations func(client *MockhttpClient, respBody *MockReadCloser)
 		wantErr      bool
 	}{
 		{
 			"successful response",
-			func(client *MockhttpClient) {
+			func(client *MockhttpClient, respBody *MockReadCloser) {
 				reqMatcher := &mockutils.HTTPRequestMatcher{
 					ExpectedRequest: newRequest(t, []byte("{\"dotty_status\":\"running\",\"ssh_info\":{\"port\":256}}")),
 				}
@@ -40,18 +40,19 @@ func Test_agentInfoUpdaterImpl_Update(t *testing.T) {
 		},
 		{
 			"unsuccessful response code",
-			func(client *MockhttpClient) {
+			func(client *MockhttpClient, respBody *MockReadCloser) {
 				reqMatcher := &mockutils.HTTPRequestMatcher{
 					ExpectedRequest: newRequest(t, []byte("{\"dotty_status\":\"running\",\"ssh_info\":{\"port\":256}}")),
 				}
 
-				client.EXPECT().Do(reqMatcher).Return(&http.Response{StatusCode: 404}, nil)
+				client.EXPECT().Do(reqMatcher).Return(&http.Response{StatusCode: 404, Body: respBody}, nil)
+				respBody.EXPECT().Close()
 			},
 			true,
 		},
 		{
 			"error from http client",
-			func(client *MockhttpClient) {
+			func(client *MockhttpClient, respBody *MockReadCloser) {
 				reqMatcher := &mockutils.HTTPRequestMatcher{
 					ExpectedRequest: newRequest(t, []byte("{\"dotty_status\":\"running\",\"ssh_info\":{\"port\":256}}")),
 				}
@@ -65,7 +66,8 @@ func Test_agentInfoUpdaterImpl_Update(t *testing.T) {
 			ctrl := gomock.NewController(t)
 			defer ctrl.Finish()
 			client := NewMockhttpClient(ctrl)
-			tt.expectations(client)
+			readCloser := NewMockReadCloser(ctrl)
+			tt.expectations(client, readCloser)
 			m := &agentInfoUpdaterImpl{
 				client: client,
 			}