From 6640d42aac8c0d780ae82a8eeb621765de3137a0 Mon Sep 17 00:00:00 2001
From: Danny <danny@alcion.ai>
Date: Wed, 7 Sep 2022 17:14:40 -0400
Subject: [PATCH] GC Restore: Event pipeline not working [FIX] (#780)

Bug Fix to ensure E2E testing. Test suite expanded.
---
 .../exchange/exchange_service_test.go         | 74 +++++++++++++++++++
 .../connector/exchange/service_functions.go   | 11 +--
 2 files changed, 76 insertions(+), 9 deletions(-)

diff --git a/src/internal/connector/exchange/exchange_service_test.go b/src/internal/connector/exchange/exchange_service_test.go
index d2ae66fa4..054676747 100644
--- a/src/internal/connector/exchange/exchange_service_test.go
+++ b/src/internal/connector/exchange/exchange_service_test.go
@@ -508,3 +508,77 @@ func (suite *ExchangeServiceSuite) TestGetRestoreContainer() {
 		})
 	}
 }
+
+// TestRestoreExchangeObject verifies path.Category usage for restored objects
+func (suite *ExchangeServiceSuite) TestRestoreExchangeObject() {
+	ctx := context.Background()
+	t := suite.T()
+	userID := tester.M365UserID(t)
+	service := loadService(t)
+	now := time.Now()
+	tests := []struct {
+		name        string
+		bytes       []byte
+		category    path.CategoryType
+		cleanupFunc func(graph.Service, string, string) error
+		destination func() string
+	}{
+		{
+			name:        "Test Mail",
+			bytes:       mockconnector.GetMockMessageBytes("Restore Exchange Object"),
+			category:    path.EmailCategory,
+			cleanupFunc: DeleteMailFolder,
+			destination: func() string {
+				folderName := "TestRestoreMailObject: " + common.FormatSimpleDateTime(now)
+				folder, err := CreateMailFolder(suite.es, userID, folderName)
+				require.NoError(t, err)
+
+				return *folder.GetId()
+			},
+		},
+		{
+			name:        "Test Contact",
+			bytes:       mockconnector.GetMockContactBytes("Test_Omega"),
+			category:    path.ContactsCategory,
+			cleanupFunc: DeleteContactFolder,
+			destination: func() string {
+				folderName := "TestRestoreContactObject: " + common.FormatSimpleDateTime(now)
+				folder, err := CreateContactFolder(suite.es, userID, folderName)
+				require.NoError(t, err)
+
+				return *folder.GetId()
+			},
+		},
+		{
+			name:        "Test Events",
+			bytes:       mockconnector.GetMockEventBytes("Restored Event Object"),
+			category:    path.EventsCategory,
+			cleanupFunc: DeleteCalendar,
+			destination: func() string {
+				calendarName := "TestRestoreEventObject: " + common.FormatSimpleDateTime(now)
+				calendar, err := CreateCalendar(suite.es, userID, calendarName)
+				require.NoError(t, err)
+
+				return *calendar.GetId()
+			},
+		},
+	}
+
+	for _, test := range tests {
+		suite.T().Run(test.name, func(t *testing.T) {
+			destination := test.destination()
+			err := RestoreExchangeObject(
+				ctx,
+				test.bytes,
+				test.category,
+				control.Copy,
+				service,
+				destination,
+				userID,
+			)
+			assert.NoError(t, err)
+			cleanupError := test.cleanupFunc(service, userID, destination)
+			assert.NoError(t, cleanupError)
+		})
+	}
+}
diff --git a/src/internal/connector/exchange/service_functions.go b/src/internal/connector/exchange/service_functions.go
index 2150f22d6..25033dedc 100644
--- a/src/internal/connector/exchange/service_functions.go
+++ b/src/internal/connector/exchange/service_functions.go
@@ -353,19 +353,12 @@ func RestoreExchangeObject(
 	service graph.Service,
 	destination, user string,
 ) error {
-	var setting optionIdentifier
-
-	switch category {
-	case path.EmailCategory, path.ContactsCategory:
-		setting = categoryToOptionIdentifier(category)
-	default:
-		return fmt.Errorf("type: %s not supported for exchange restore", category)
-	}
-
 	if policy != control.Copy {
 		return fmt.Errorf("restore policy: %s not supported", policy)
 	}
 
+	setting := categoryToOptionIdentifier(category)
+
 	switch setting {
 	case messages:
 		return RestoreMailMessage(ctx, bits, service, control.Copy, destination, user)