B-21669 - fix tests and error handling on endpoint.

transcom · Jan 8, 2025 · 2fa55a1 · 2fa55a1
1 parent 6e122f4
commit 2fa55a1
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 34 deletions.
diff --git a/pkg/handlers/internalapi/uploads.go b/pkg/handlers/internalapi/uploads.go
@@ -259,10 +259,11 @@ type GetUploadStatusHandler struct {
 }
 
 type CustomNewUploadStatusOK struct {
-	params   uploadop.GetUploadStatusParams
-	appCtx   appcontext.AppContext
-	receiver notifications.NotificationReceiver
-	storer   storage.FileStorer
+	params     uploadop.GetUploadStatusParams
+	storageKey string
+	appCtx     appcontext.AppContext
+	receiver   notifications.NotificationReceiver
+	storer     storage.FileStorer
 }
 
 // AVStatusType represents the type of the anti-virus status, whether it is still processing, clean or infected
@@ -289,23 +290,8 @@ func writeEventStreamMessage(rw http.ResponseWriter, producer runtime.Producer,
 
 func (o *CustomNewUploadStatusOK) WriteResponse(rw http.ResponseWriter, producer runtime.Producer) {
 
-	// TODO: add check for permissions to view upload
-
-	uploadId := o.params.UploadID.String()
-
-	uploadUUID, err := uuid.FromString(uploadId)
-	if err != nil {
-		uploadop.NewGetUploadStatusInternalServerError().WriteResponse(rw, producer)
-		panic(err)
-	}
-
 	// Check current tag before event-driven wait for anti-virus
-	uploaded, err := models.FetchUserUploadFromUploadID(o.appCtx.DB(), o.appCtx.Session(), uploadUUID)
-	if err != nil {
-		panic(err)
-	}
-
-	tags, err := o.storer.Tags(uploaded.Upload.StorageKey)
+	tags, err := o.storer.Tags(o.storageKey)
 	var uploadStatus AVStatusType
 	if err != nil || len(tags) == 0 {
 		uploadStatus = AVStatusTypePROCESSING
@@ -335,7 +321,7 @@ func (o *CustomNewUploadStatusOK) WriteResponse(rw http.ResponseWriter, producer
 					]
 				}
 			}
-	}`, uploadId)
+	}`, o.params.UploadID)
 
 	notificationParams := notifications.NotificationQueueParams{
 		SubscriptionTopicName: topicName,
@@ -370,7 +356,7 @@ func (o *CustomNewUploadStatusOK) WriteResponse(rw http.ResponseWriter, producer
 		if len(messages) != 0 {
 			errTransaction := o.appCtx.NewTransaction(func(txnAppCtx appcontext.AppContext) error {
 
-				tags, err := o.storer.Tags(uploaded.Upload.StorageKey)
+				tags, err := o.storer.Tags(o.storageKey)
 
 				if err != nil || len(tags) == 0 {
 					uploadStatus = AVStatusTypePROCESSING
@@ -405,11 +391,36 @@ func (o *CustomNewUploadStatusOK) WriteResponse(rw http.ResponseWriter, producer
 func (h GetUploadStatusHandler) Handle(params uploadop.GetUploadStatusParams) middleware.Responder {
 	return h.AuditableAppContextFromRequestWithErrors(params.HTTPRequest,
 		func(appCtx appcontext.AppContext) (middleware.Responder, error) {
+
+			handleError := func(err error) (middleware.Responder, error) {
+				appCtx.Logger().Error("GetUploadStatusHandler error", zap.Error(err))
+				switch errors.Cause(err) {
+				case models.ErrFetchForbidden:
+					return uploadop.NewGetUploadStatusForbidden(), err
+				case models.ErrFetchNotFound:
+					return uploadop.NewGetUploadStatusNotFound(), err
+				default:
+					return uploadop.NewGetUploadStatusInternalServerError(), err
+				}
+			}
+
+			uploadId := params.UploadID.String()
+			uploadUUID, err := uuid.FromString(uploadId)
+			if err != nil {
+				return handleError(err)
+			}
+
+			uploaded, err := models.FetchUserUploadFromUploadID(appCtx.DB(), appCtx.Session(), uploadUUID)
+			if err != nil {
+				return handleError(err)
+			}
+
 			return &CustomNewUploadStatusOK{
-				params:   params,
-				appCtx:   h.AppContextFromRequest(params.HTTPRequest),
-				receiver: h.NotificationReceiver(),
-				storer:   h.FileStorer(),
+				params:     params,
+				storageKey: uploaded.Upload.StorageKey,
+				appCtx:     h.AppContextFromRequest(params.HTTPRequest),
+				receiver:   h.NotificationReceiver(),
+				storer:     h.FileStorer(),
 			}, nil
 		})
 }

diff --git a/pkg/handlers/internalapi/uploads_test.go b/pkg/handlers/internalapi/uploads_test.go
@@ -449,15 +449,14 @@ func (suite *HandlerSuite) TestDeleteUploadHandlerSuccessEvenWithS3Failure() {
 	suite.NotNil(queriedUpload.DeletedAt)
 }
 
-// TODO: functioning test
 func (suite *HandlerSuite) TestGetUploadStatusHandlerSuccess() {
 	fakeS3 := storageTest.NewFakeS3Storage(true)
 	localReceiver := notifications.StubNotificationReceiver{}
 
-	move := factory.BuildMove(suite.DB(), nil, nil)
+	orders := factory.BuildOrder(suite.DB(), nil, nil)
 	uploadUser1 := factory.BuildUserUpload(suite.DB(), []factory.Customization{
 		{
-			Model:    move.Orders.UploadedOrders,
+			Model:    orders.UploadedOrders,
 			LinkOnly: true,
 		},
 		{
@@ -470,10 +469,11 @@ func (suite *HandlerSuite) TestGetUploadStatusHandlerSuccess() {
 	}, nil)
 
 	file := suite.Fixture(FixturePDF)
-	fakeS3.Store(uploadUser1.Upload.StorageKey, file.Data, "somehash", nil)
+	_, err := fakeS3.Store(uploadUser1.Upload.StorageKey, file.Data, "somehash", nil)
+	suite.NoError(err)
 
 	params := uploadop.NewGetUploadStatusParams()
-	params.UploadID = strfmt.UUID(uploadUser1.ID.String())
+	params.UploadID = strfmt.UUID(uploadUser1.Upload.ID.String())
 
 	req := &http.Request{}
 	req = suite.AuthenticateRequest(req, uploadUser1.Document.ServiceMember)
@@ -490,8 +490,42 @@ func (suite *HandlerSuite) TestGetUploadStatusHandlerSuccess() {
 	suite.True(ok)
 
 	queriedUpload := models.Upload{}
-	err := suite.DB().Find(&queriedUpload, uploadUser1.Upload.ID)
-	suite.Nil(err)
+	err = suite.DB().Find(&queriedUpload, uploadUser1.Upload.ID)
+	suite.NoError(err)
+}
+
+func (suite *HandlerSuite) TestGetUploadStatusHandlerFailure() {
+	suite.Run("Error on no match for uploadId", func() {
+		orders := factory.BuildOrder(suite.DB(), factory.GetTraitActiveServiceMemberUser(), nil)
+
+		uploadUUID := uuid.Must(uuid.NewV4())
+
+		params := uploadop.NewGetUploadStatusParams()
+		params.UploadID = strfmt.UUID(uploadUUID.String())
+
+		req := &http.Request{}
+		req = suite.AuthenticateRequest(req, orders.ServiceMember)
+		params.HTTPRequest = req
+
+		fakeS3 := storageTest.NewFakeS3Storage(true)
+		localReceiver := notifications.StubNotificationReceiver{}
+
+		handlerConfig := suite.HandlerConfig()
+		handlerConfig.SetFileStorer(fakeS3)
+		handlerConfig.SetNotificationReceiver(localReceiver)
+		uploadInformationFetcher := upload.NewUploadInformationFetcher()
+		handler := GetUploadStatusHandler{handlerConfig, uploadInformationFetcher}
+
+		response := handler.Handle(params)
+		_, ok := response.(*uploadop.GetUploadStatusNotFound)
+		suite.True(ok)
+
+		queriedUpload := models.Upload{}
+		err := suite.DB().Find(&queriedUpload, uploadUUID)
+		suite.Error(err)
+	})
+
+	// TODO: ADD A FORBIDDEN TEST
 }
 
 func (suite *HandlerSuite) TestCreatePPMUploadsHandlerSuccess() {

diff --git a/pkg/handlers/routing/internalapi_test/uploads_test.go b/pkg/handlers/routing/internalapi_test/uploads_test.go
@@ -36,6 +36,6 @@ func (suite *InternalAPISuite) TestUploads() {
 
 		suite.Equal(http.StatusOK, rr.Code)
 		suite.Equal("text/event-stream", rr.Header().Get("content-type"))
-		suite.Equal("id: 0\nevent: message\ndata: CLEAN\n\n", rr.Body.String())
+		suite.Equal("id: 0\nevent: message\ndata: CLEAN\n\nid: 1\nevent: close\ndata: Connection closed\n\n", rr.Body.String())
 	})
 }