From 61f769cc15239369bfbcb8d56b71adfd7bbdc51b Mon Sep 17 00:00:00 2001
From: Danny <danny@alcion.ai>
Date: Tue, 13 Sep 2022 11:33:08 -0400
Subject: [PATCH] 812 gc retrieval function tests (#817)

Issue 812: Test Expansion
Retrieval functions added to the test exchange package.

Issue #831:
Retry handler added to message restore pipeline.
---
 .../exchange/exchange_service_test.go         | 65 +++++++++++++++++++
 .../connector/exchange/service_functions.go   | 15 ++++-
 .../connector/exchange/service_query.go       |  2 +-
 3 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/src/internal/connector/exchange/exchange_service_test.go b/src/internal/connector/exchange/exchange_service_test.go
index 9b4ab0140..b9059c006 100644
--- a/src/internal/connector/exchange/exchange_service_test.go
+++ b/src/internal/connector/exchange/exchange_service_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/microsoftgraph/msgraph-sdk-go/models"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
@@ -347,6 +348,70 @@ func (suite *ExchangeServiceSuite) TestParseCalendarIDFromEvent() {
 	}
 }
 
+// TestRetrievalFunctions ensures that utility functions used
+// to transform work within the current version of GraphAPI.
+func (suite *ExchangeServiceSuite) TestRetrievalFunctions() {
+	var (
+		userID   = tester.M365UserID(suite.T())
+		objectID string
+	)
+
+	tests := []struct {
+		name         string
+		query        GraphQuery
+		retrieveFunc GraphRetrievalFunc
+	}{
+		{
+			name:         "Test Retrieve Message Function",
+			query:        GetAllMessagesForUser,
+			retrieveFunc: RetrieveMessageDataForUser,
+		},
+		{
+			name:         "Test Retrieve Contact Function",
+			query:        GetAllContactsForUser,
+			retrieveFunc: RetrieveContactDataForUser,
+		},
+		{
+			name:         "Test Retrieve Event Function",
+			query:        GetAllEventsForUser,
+			retrieveFunc: RetrieveEventDataForUser,
+		},
+	}
+
+	for _, test := range tests {
+		suite.T().Run(test.name, func(t *testing.T) {
+			output, err := test.query(suite.es, userID)
+			require.NoError(t, err)
+			switch v := output.(type) {
+			case *models.MessageCollectionResponse:
+				transform := output.(models.MessageCollectionResponseable)
+				response := transform.GetValue()
+				require.Greater(t, len(response), 0)
+
+				objectID = *response[0].GetId()
+			case *models.ContactCollectionResponse:
+				transform := output.(models.ContactCollectionResponseable)
+				response := transform.GetValue()
+				require.Greater(t, len(response), 0)
+
+				objectID = *response[0].GetId()
+			case *models.EventCollectionResponse:
+				transform := output.(models.EventCollectionResponseable)
+				response := transform.GetValue()
+				require.Greater(t, len(response), 0)
+
+				objectID = *response[0].GetId()
+			default:
+				t.Logf("What is this type: %T\n", v)
+			}
+			require.NotEmpty(t, objectID)
+			retrieved, err := test.retrieveFunc(suite.es, userID, objectID)
+			assert.NoError(t, err, support.ConnectorStackErrorTrace(err))
+			assert.NotNil(t, retrieved)
+		})
+	}
+}
+
 // TestGetMailFolderID verifies the ability to retrieve folder ID of folders
 // at the top level of the file tree
 func (suite *ExchangeServiceSuite) TestGetContainerID() {
diff --git a/src/internal/connector/exchange/service_functions.go b/src/internal/connector/exchange/service_functions.go
index 50dc41ed5..a60ce2a1c 100644
--- a/src/internal/connector/exchange/service_functions.go
+++ b/src/internal/connector/exchange/service_functions.go
@@ -520,7 +520,7 @@ func RestoreMailMessage(
 	// Creates messageable object from original bytes
 	originalMessage, err := support.CreateMessageFromBytes(bits)
 	if err != nil {
-		return err
+		return errors.Wrapf(err, "restore mail message rcvd: %v", bits)
 	}
 	// Sets fields from original message from storage
 	clone := support.ToMessage(originalMessage)
@@ -566,7 +566,18 @@ func RestoreMailMessage(
 // @param destination represents M365 ID of a folder within the users's space
 // @param message is a models.Messageable interface from "github.com/microsoftgraph/msgraph-sdk-go/models"
 func SendMailToBackStore(service graph.Service, user, destination string, message models.Messageable) error {
-	sentMessage, err := service.Client().UsersById(user).MailFoldersById(destination).Messages().Post(message)
+	var (
+		sentMessage models.Messageable
+		err         error
+	)
+
+	for count := 0; count < numberOfRetries; count++ {
+		sentMessage, err = service.Client().UsersById(user).MailFoldersById(destination).Messages().Post(message)
+		if err == nil && sentMessage != nil {
+			break
+		}
+	}
+
 	if err != nil {
 		return support.WrapAndAppend(": "+support.ConnectorStackErrorTrace(err), err, nil)
 	}
diff --git a/src/internal/connector/exchange/service_query.go b/src/internal/connector/exchange/service_query.go
index ac0b05b2a..3ba60cbb2 100644
--- a/src/internal/connector/exchange/service_query.go
+++ b/src/internal/connector/exchange/service_query.go
@@ -17,7 +17,7 @@ import (
 // into M365 backstore. Responses -> returned items will only contain the information
 // that is included in the options
 // TODO: use selector or path for granularity into specific folders or specific date ranges
-type GraphQuery func(graph.Service, string) (absser.Parsable, error)
+type GraphQuery func(gs graph.Service, userID string) (absser.Parsable, error)
 
 // GetAllMessagesForUser is a GraphQuery function for receiving all messages for a single user
 func GetAllMessagesForUser(gs graph.Service, user string) (absser.Parsable, error) {