From f0ff7ee982d59a441e6bf85cda4890426e04797c Mon Sep 17 00:00:00 2001
From: ashmrtn <ashmrtnz@alcion.ai>
Date: Wed, 28 Sep 2022 15:41:50 -0700
Subject: [PATCH] Basic checks for bad/empty values for restore (#973)

## Description

Test bad services or empty collections don't cause an error

## Type of change

<!--- Please check the type of change your PR introduces: --->
- [ ] :sunflower: Feature
- [ ] :bug: Bugfix
- [ ] :world_map: Documentation
- [x] :robot: Test
- [ ] :computer: CI/Deployment
- [ ] :hamster: Trivial/Minor

## Issue(s)

* #913

## Test Plan

<!-- How will this be tested prior to merging.-->
- [ ] :muscle: Manual
- [x] :zap: Unit test
- [ ] :green_heart: E2E
---
 src/internal/connector/graph_connector.go     |  7 ++-
 .../graph_connector_disconnected_test.go      | 20 +++++++
 .../connector/graph_connector_test.go         | 52 +++++++++++++++++++
 3 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/internal/connector/graph_connector.go b/src/internal/connector/graph_connector.go
index 27ed21c73..e736af0aa 100644
--- a/src/internal/connector/graph_connector.go
+++ b/src/internal/connector/graph_connector.go
@@ -346,10 +346,15 @@ func (gc *GraphConnector) AwaitStatus() *support.ConnectorOperationStatus {
 
 // UpdateStatus is used by gc initiated tasks to indicate completion
 func (gc *GraphConnector) UpdateStatus(status *support.ConnectorOperationStatus) {
+	defer gc.wg.Done()
+
+	if status == nil {
+		return
+	}
+
 	gc.mu.Lock()
 	defer gc.mu.Unlock()
 	gc.status = support.MergeStatus(gc.status, *status)
-	gc.wg.Done()
 }
 
 // Status returns the current status of the graphConnector operaion.
diff --git a/src/internal/connector/graph_connector_disconnected_test.go b/src/internal/connector/graph_connector_disconnected_test.go
index 9f39cb921..118e15537 100644
--- a/src/internal/connector/graph_connector_disconnected_test.go
+++ b/src/internal/connector/graph_connector_disconnected_test.go
@@ -10,10 +10,13 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 
+	"github.com/alcionai/corso/src/internal/common"
 	"github.com/alcionai/corso/src/internal/connector/support"
 	"github.com/alcionai/corso/src/internal/tester"
 	"github.com/alcionai/corso/src/pkg/account"
+	"github.com/alcionai/corso/src/pkg/control"
 	"github.com/alcionai/corso/src/pkg/credentials"
+	"github.com/alcionai/corso/src/pkg/selectors"
 )
 
 // ---------------------------------------------------------------
@@ -172,3 +175,20 @@ func (suite *DisconnectedGraphConnectorSuite) TestGraphConnector_ErrorChecking()
 		})
 	}
 }
+
+func (suite *DisconnectedGraphConnectorSuite) TestRestoreFailsBadService() {
+	t := suite.T()
+	ctx := context.Background()
+	gc := GraphConnector{wg: &sync.WaitGroup{}}
+	sel := selectors.Selector{
+		Service: selectors.ServiceUnknown,
+	}
+	dest := control.DefaultRestoreDestination(common.SimpleDateTimeFormatOneDrive)
+
+	assert.Error(t, gc.RestoreDataCollections(ctx, sel, dest, nil))
+
+	status := gc.AwaitStatus()
+	assert.Equal(t, 0, status.ObjectCount)
+	assert.Equal(t, 0, status.FolderCount)
+	assert.Equal(t, 0, status.Successful)
+}
diff --git a/src/internal/connector/graph_connector_test.go b/src/internal/connector/graph_connector_test.go
index b26b86d4e..dfad37d8a 100644
--- a/src/internal/connector/graph_connector_test.go
+++ b/src/internal/connector/graph_connector_test.go
@@ -403,6 +403,58 @@ func (suite *GraphConnectorIntegrationSuite) TestRestoreContact() {
 	suite.T().Log(value.String())
 }
 
+func (suite *GraphConnectorIntegrationSuite) TestEmptyCollections() {
+	dest := control.DefaultRestoreDestination(common.SimpleDateTimeFormatOneDrive)
+	table := []struct {
+		name string
+		col  []data.Collection
+		sel  selectors.Selector
+	}{
+		{
+			name: "ExchangeNil",
+			col:  nil,
+			sel: selectors.Selector{
+				Service: selectors.ServiceExchange,
+			},
+		},
+		{
+			name: "ExchangeEmpty",
+			col:  []data.Collection{},
+			sel: selectors.Selector{
+				Service: selectors.ServiceExchange,
+			},
+		},
+		{
+			name: "OneDriveNil",
+			col:  nil,
+			sel: selectors.Selector{
+				Service: selectors.ServiceOneDrive,
+			},
+		},
+		{
+			name: "OneDriveEmpty",
+			col:  []data.Collection{},
+			sel: selectors.Selector{
+				Service: selectors.ServiceOneDrive,
+			},
+		},
+	}
+
+	for _, test := range table {
+		suite.T().Run(test.name, func(t *testing.T) {
+			ctx := context.Background()
+
+			err := suite.connector.RestoreDataCollections(ctx, test.sel, dest, test.col)
+			require.NoError(t, err)
+
+			stats := suite.connector.AwaitStatus()
+			assert.Zero(t, stats.ObjectCount)
+			assert.Zero(t, stats.FolderCount)
+			assert.Zero(t, stats.Successful)
+		})
+	}
+}
+
 func (suite *GraphConnectorIntegrationSuite) TestRestoreAndBackup() {
 	bodyText := "This email has some text. However, all the text is on the same line."
 	subjectText := "Test message for restore"