fix(digestwriter): accept different account number types

Jakub007d · Sep 4, 2023 · e5ae135 · e5ae135
1 parent 206d813
commit e5ae135
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 20 deletions.
diff --git a/base/utils/payload_tracker.go b/base/utils/payload_tracker.go
@@ -2,7 +2,6 @@ package utils
 
 import (
 	"encoding/json"
-	"strconv"
 	"time"
 
 	"github.com/Shopify/sarama"
@@ -28,8 +27,8 @@ type PayloadTrackerEvent struct {
 	Date        *string `json:"date"` // RFC3339
 }
 
-func (e *PayloadTrackerEvent) SetOrgIDFromUint(id uint32) {
-	orgID := strconv.FormatUint(uint64(id), 10)
+func (e *PayloadTrackerEvent) SetOrgIDFromUint(id json.Number) {
+	orgID := id.String()
 	e.OrgID = &orgID
 }
 

diff --git a/digestwriter/consumer.go b/digestwriter/consumer.go
@@ -38,12 +38,6 @@ const (
 	errClusterData = "error updating cluster data"
 )
 
-// OrgID data type represents organization ID.
-type OrgID uint32
-
-// AccountNumber data type represents account number for a given report.
-type AccountNumber uint32
-
 // ClusterName represents the external cluster UUID contained in the consumed message
 type ClusterName string
 
@@ -77,13 +71,13 @@ type Namespace struct {
 // IncomingMessage data structure is representation of message consumed from
 // the configured topic
 type IncomingMessage struct {
-	Organization  OrgID         `json:"OrgID"`
-	AccountNumber AccountNumber `json:"AccountNumber"`
-	ClusterName   ClusterName   `json:"ClusterName"`
-	Workload      *Workload     `json:"Images"`
-	LastChecked   string        `json:"-"`
-	Version       uint8         `json:"Version"`
-	RequestID     RequestID     `json:"RequestID"`
+	Organization  json.Number `json:"OrgID"`
+	AccountNumber json.Number `json:"AccountNumber"`
+	ClusterName   ClusterName `json:"ClusterName"`
+	Workload      *Workload   `json:"Images"`
+	LastChecked   string      `json:"-"`
+	Version       uint8       `json:"Version"`
+	RequestID     RequestID   `json:"RequestID"`
 }
 
 // DigestConsumer Struct that must fulfill the Processor interface defined in utils/kafka.go
@@ -166,7 +160,7 @@ func (d *DigestConsumer) ProcessMessage(msg *sarama.ConsumerMessage) error {
 
 	// Set up payload tracker event
 	ptEvent := utils.NewPayloadTrackerEvent(string(message.RequestID))
-	ptEvent.SetOrgIDFromUint(uint32(message.Organization))
+	ptEvent.SetOrgIDFromUint(message.Organization)
 
 	// Send Payload Tracker message with status received
 	ptEvent.UpdateStatusReceived()
@@ -264,6 +258,9 @@ func parseMessage(messageValue []byte) (IncomingMessage, error) {
 	if deserialized.Workload == nil {
 		return deserialized, errors.New("missing required attribute 'Images'")
 	}
+	if deserialized.Organization.String() == "" {
+		return deserialized, errors.New("OrgID cannot be null")
+	}
 
 	logger.WithFields(logrus.Fields{
 		requestIDKey: deserialized.RequestID,

diff --git a/digestwriter/consumer_test.go b/digestwriter/consumer_test.go
@@ -464,3 +464,29 @@ func TestProcessMessageUUID(t *testing.T) {
 	assert.Equal(t, "error", ptEvent.Status)
 	assert.Equal(t, "error updating cluster data", ptEvent.StatusMsg)
 }
+
+func TestDigestMessageParse(t *testing.T) {
+	SetupLogger()
+	utils.SetupLogger()
+
+	validCases := [][]byte{
+		[]byte(`{"OrgID":12341446,"AccountNumber":6341839,"ClusterName":"04a816ea-cd0b-47c3-b754-a9008b127d84","Images":{"imageCount":2,"images":{},"namespaces":{}},"Version":2,"RequestID":"cbcbfeb72f074dffad1528dd209b130e"}`),
+		[]byte(`{"OrgID":"12341446","AccountNumber":6341839,"ClusterName":"04a816ea-cd0b-47c3-b754-a9008b127d84","Images":{"imageCount":2,"images":{},"namespaces":{}},"Version":2,"RequestID":"cbcbfeb72f074dffad1528dd209b130e"}`),
+		[]byte(`{"OrgID":12341446,"AccountNumber":"6341839","ClusterName":"04a816ea-cd0b-47c3-b754-a9008b127d84","Images":{"imageCount":2,"images":{},"namespaces":{}},"Version":2,"RequestID":"cbcbfeb72f074dffad1528dd209b130e"}`),
+		[]byte(`{"OrgID":12341446,"AccountNumber":null,"ClusterName":"04a816ea-cd0b-47c3-b754-a9008b127d84","Images":{"imageCount":2,"images":{},"namespaces":{}},"Version":2,"RequestID":"cbcbfeb72f074dffad1528dd209b130e"}`),
+	}
+	invalidCases := [][]byte{
+		[]byte(`{"OrgID":null,"AccountNumber":6341839,"ClusterName":"04a816ea-cd0b-47c3-b754-a9008b127d84","Images":{"imageCount":2,"images":{},"namespaces":{}},"Version":2,"RequestID":"cbcbfeb72f074dffad1528dd209b130e"}`),
+		[]byte(`{"AccountNumber":6341839,"ClusterName":"04a816ea-cd0b-47c3-b754-a9008b127d84","Images":{"imageCount":2,"images":{},"namespaces":{}},"Version":2,"RequestID":"cbcbfeb72f074dffad1528dd209b130e"}`),
+	}
+
+	for _, msg := range validCases {
+		_, err := parseMessage(msg)
+		assert.Nil(t, err)
+	}
+
+	for _, msg := range invalidCases {
+		_, err := parseMessage(msg)
+		assert.Equal(t, "OrgID cannot be null", err.Error())
+	}
+}
diff --git a/digestwriter/storage.go b/digestwriter/storage.go
@@ -6,6 +6,7 @@ package digestwriter
 import (
 	"app/base/models"
 	"app/base/utils"
+	"encoding/json"
 	"fmt"
 	"time"
 
@@ -21,7 +22,7 @@ const (
 
 // Storage represents an interface to almost any database or storage system
 type Storage interface {
-	WriteClusterInfo(cluster ClusterName, orgID OrgID, workload Workload, digests []string) error
+	WriteClusterInfo(cluster ClusterName, orgID json.Number, workload Workload, digests []string) error
 }
 
 // DBStorage is an implementation of Storage
@@ -189,7 +190,7 @@ func (storage *DBStorage) linkDigestsToCluster(tx *gorm.DB, clusterStr string, c
 }
 
 // WriteClusterInfo updates the 'cluster' table with the provided info
-func (storage *DBStorage) WriteClusterInfo(cluster ClusterName, orgID OrgID, workload Workload, digests []string) error {
+func (storage *DBStorage) WriteClusterInfo(cluster ClusterName, orgID json.Number, workload Workload, digests []string) error {
 	// prepare data
 	clusterStr := string(cluster)
 	clusterUUID, err := uuid.Parse(clusterStr)

diff --git a/digestwriter/storage_test.go b/digestwriter/storage_test.go
@@ -7,6 +7,7 @@ import (
 	"app/base/models"
 	"app/digestwriter"
 	"database/sql"
+	"encoding/json"
 	"regexp"
 	"testing"
 	"time"
@@ -22,7 +23,7 @@ var (
 	digestModifiedAtTime = time.Now().UTC()
 	clusterName          = digestwriter.ClusterName("84f7eedc-0000-0000-9d4d-000000000000")
 	invalidClusterName   = digestwriter.ClusterName("99z7zzzz-0000-0000-9d4d-000000000000")
-	testOrgID            = digestwriter.OrgID(1)
+	testOrgID            = json.Number("1")
 	workload             = digestwriter.Workload{}
 
 	anyArgForMockSQLQueries = sqlmock.AnyArg()