Fix CountQuery paramter-parsing

Author: Jaskaranbir

inventory/query_count.go

@@ -3,7 +3,6 @@ package inventory
 import (
 	"encoding/json"
 	"log"
-	"strconv"
 
 	"github.com/mongodb/mongo-go-driver/mongo/findopt"
 
@@ -13,8 +12,27 @@ import (
 	"github.com/TerrexTech/go-mongoutils/mongo"
 )
 
+type CountQuery struct {
+	Count int `json:"count,omitempty"`
+}
+
 func queryCount(collection *mongo.Collection, event *model.Event) *model.KafkaResponse {
-	count, err := strconv.Atoi(string(event.Data))
+	countQuery := &CountQuery{}
+	err := json.Unmarshal(event.Data, countQuery)
+	if err != nil {
+		err = errors.Wrap(err, "QueryCount: Error unmarshalling paramters")
+		log.Println(err)
+		return &model.KafkaResponse{
+			AggregateID:   event.AggregateID,
+			CorrelationID: event.CorrelationID,
+			Error:         err.Error(),
+			ErrorCode:     InternalError,
+			EventAction:   event.EventAction,
+			ServiceAction: event.ServiceAction,
+			UUID:          event.UUID,
+		}
+	}
+
 	if err != nil {
 		err = errors.Wrap(err, "QueryCount: Error converting parameter to number")
 		log.Println(err)
@@ -29,7 +47,7 @@ func queryCount(collection *mongo.Collection, event *model.Event) *model.KafkaRe
 		}
 	}
 
-	if count == 0 {
+	if countQuery.Count == 0 {
 		err = errors.New("count must be greater than 0")
 		log.Println(err)
 		return &model.KafkaResponse{
@@ -42,7 +60,7 @@ func queryCount(collection *mongo.Collection, event *model.Event) *model.KafkaRe
 			UUID:          event.UUID,
 		}
 	}
-	if count > 100 {
+	if countQuery.Count > 100 {
 		err = errors.New("count must be less than 100")
 		log.Println(err)
 		return &model.KafkaResponse{
@@ -77,6 +95,6 @@ func queryCount(collection *mongo.Collection, event *model.Event) *model.KafkaRe
 	}
 	event.Data = eventData
 
-	findopts := findopt.Limit(int64(count))
+	findopts := findopt.Limit(int64(countQuery.Count))
 	return queryInventory(collection, event, findopts)
 }

test/inventory_suite_test.go

@@ -290,7 +290,10 @@ var _ = Describe("InventoryAggregate", func() {
 
 			Byf("Creating query args")
 			// Get 3 records
-			marshalQuery, err := json.Marshal(3)
+			query := &inventory.CountQuery{
+				Count: 2,
+			}
+			marshalQuery, err := json.Marshal(query)
 			Expect(err).ToNot(HaveOccurred())
 
 			Byf("Creating query MockEvent")
@@ -335,7 +338,6 @@ var _ = Describe("InventoryAggregate", func() {
 					err = json.Unmarshal(kr.Result, &result)
 					Expect(err).ToNot(HaveOccurred())
 					Expect(result).ToNot(BeEmpty())
-					log.Println(string(kr.Result))
 					return true
 				}
 				return false