From c7a7dde40eabc25ec7dc49e1d2d7c075a166c566 Mon Sep 17 00:00:00 2001
From: udsamani <uditpsamani@hotmail.com>
Date: Wed, 11 Sep 2024 02:15:25 +0100
Subject: [PATCH] Fix failing tests

---
 pkg/jobstore/boltdb/store.go              |  2 +-
 pkg/orchestrator/endpoint.go              |  4 ++--
 pkg/publicapi/client/v2/client.go         | 29 +++++++++++++++++++----
 pkg/publicapi/middleware/error_handler.go |  7 ++++++
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/pkg/jobstore/boltdb/store.go b/pkg/jobstore/boltdb/store.go
index 183c8ab2ba..7a9836a585 100644
--- a/pkg/jobstore/boltdb/store.go
+++ b/pkg/jobstore/boltdb/store.go
@@ -760,7 +760,7 @@ func (b *BoltJobStore) CreateJob(ctx context.Context, job models.Job) error {
 	job.Normalize()
 	err := job.Validate()
 	if err != nil {
-		return err
+		return jobstore.NewJobStoreError(err.Error())
 	}
 	return b.update(ctx, func(tx *bolt.Tx) (err error) {
 		return b.createJob(tx, job)
diff --git a/pkg/orchestrator/endpoint.go b/pkg/orchestrator/endpoint.go
index 60e5f7d17a..1d9893d7ea 100644
--- a/pkg/orchestrator/endpoint.go
+++ b/pkg/orchestrator/endpoint.go
@@ -148,14 +148,14 @@ func (e *BaseEndpoint) StopJob(ctx context.Context, request *StopJobRequest) (St
 		// no need to stop a job that is already stopped
 		return StopJobResponse{}, nil
 	case models.JobStateTypeCompleted:
-		return StopJobResponse{}, fmt.Errorf("cannot stop job in state %s", job.State.StateType)
+		return StopJobResponse{}, models.NewBaseError("cannot stop job in state %s", job.State.StateType)
 	default:
 		// continue
 	}
 
 	txContext, err := e.store.BeginTx(ctx)
 	if err != nil {
-		return StopJobResponse{}, fmt.Errorf("failed to begin transaction: %w", err)
+		return StopJobResponse{}, jobstore.NewBoltDbError(err.Error())
 	}
 
 	defer func() {
diff --git a/pkg/publicapi/client/v2/client.go b/pkg/publicapi/client/v2/client.go
index db68b835ad..6e0b9d377d 100644
--- a/pkg/publicapi/client/v2/client.go
+++ b/pkg/publicapi/client/v2/client.go
@@ -57,6 +57,9 @@ func (c *httpClient) Get(ctx context.Context, endpoint string, in apimodels.GetR
 	r := in.ToHTTPRequest()
 
 	_, resp, err := c.doRequest(ctx, http.MethodGet, endpoint, r) //nolint:bodyclose // this is being closed
+	if err != nil {
+		return err
+	}
 
 	if resp.StatusCode == http.StatusUnauthorized {
 		return apimodels.NewUnauthorizedError("invalid token")
@@ -93,13 +96,28 @@ func (c *httpClient) write(ctx context.Context, verb, endpoint string, in apimod
 	if r.BodyObj == nil && r.Body == nil {
 		r.BodyObj = in
 	}
-	_, resp, err := requireOK(c.doRequest(ctx, verb, endpoint, r)) //nolint:bodyclose // this is being closed
-	if err != nil && resp != nil && resp.StatusCode == http.StatusUnauthorized {
-		return apimodels.ErrInvalidToken
-	} else if err != nil {
+
+	_, resp, err := c.doRequest(ctx, verb, endpoint, r) //nolint:bodyclose // this is being closed
+	defer resp.Body.Close()
+	if err != nil {
 		return err
 	}
-	defer resp.Body.Close()
+
+	if resp.StatusCode == http.StatusUnauthorized {
+		return apimodels.ErrInvalidToken
+	}
+
+	var apiError *apimodels.APIError
+	if resp.StatusCode != http.StatusOK {
+		apiError, err = apimodels.FromHttpResponse(resp)
+		if err != nil {
+			return err
+		}
+	}
+
+	if apiError != nil {
+		return apiError
+	}
 
 	if out != nil {
 		if err := decodeBody(resp, &out); err != nil {
@@ -107,6 +125,7 @@ func (c *httpClient) write(ctx context.Context, verb, endpoint string, in apimod
 		}
 		out.Normalize()
 	}
+
 	return nil
 }
 
diff --git a/pkg/publicapi/middleware/error_handler.go b/pkg/publicapi/middleware/error_handler.go
index 7b454d8a3d..542462358a 100644
--- a/pkg/publicapi/middleware/error_handler.go
+++ b/pkg/publicapi/middleware/error_handler.go
@@ -21,11 +21,18 @@ func CustomHTTPErrorHandler(err error, c echo.Context) {
 		code = e.Code().HTTPStatusCode()
 		message = e.Error()
 
+	case *echo.HTTPError:
+		// This is needed, in case any other middleware throws an error. In
+		// such a scenario we just use it as the error code.
+		code = e.Code
+		message = e.Message.(string)
+
 	default:
 		// In an ideal world this should never happen. We should always have are errors
 		// from server as APIError. If output is this generic string, one should evaluate
 		// and map it to APIError and send in appropriate message.= http.StatusInternalServerError
 		message = "internal server error"
+		code = c.Response().Status
 	}
 
 	requestID := c.Request().Header.Get(echo.HeaderXRequestID)