From b592ef4749e5d12f8847e5e94bc14aa42d6c7a59 Mon Sep 17 00:00:00 2001
From: Abin Simon <mail@meain.io>
Date: Tue, 17 Oct 2023 14:34:17 +0530
Subject: [PATCH] Fix failing tests

---
 src/internal/m365/controller.go | 28 ++++++++++++++++++++++++++++
 src/pkg/control/restore_test.go | 23 ++++++++++++++---------
 2 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/src/internal/m365/controller.go b/src/internal/m365/controller.go
index 929300d7f..82b2705c3 100644
--- a/src/internal/m365/controller.go
+++ b/src/internal/m365/controller.go
@@ -10,6 +10,8 @@ import (
 	"github.com/alcionai/corso/src/internal/common/idname"
 	"github.com/alcionai/corso/src/internal/data"
 	"github.com/alcionai/corso/src/internal/m365/graph"
+	"github.com/alcionai/corso/src/internal/m365/resource"
+	"github.com/alcionai/corso/src/internal/m365/service/common"
 	"github.com/alcionai/corso/src/internal/m365/support"
 	"github.com/alcionai/corso/src/internal/operations/inject"
 	"github.com/alcionai/corso/src/pkg/account"
@@ -80,9 +82,35 @@ func NewController(
 		return nil, clues.Wrap(err, "creating api client").WithClues(ctx)
 	}
 
+	var rCli *common.ResourceClient
+
+	// no failure for unknown service.
+	// In that case we create a controller that doesn't attempt to look up any resource
+	// data.  This case helps avoid unnecessary service calls when the end user is running
+	// repo init and connect commands via the CLI.  All other callers should be expected
+	// to pass in a known service, or else expect downstream failures.
+	if pst != path.UnknownService {
+		rc := resource.UnknownResource
+
+		switch pst {
+		case path.ExchangeService, path.OneDriveService:
+			rc = resource.Users
+		case path.GroupsService:
+			rc = resource.Groups
+		case path.SharePointService:
+			rc = resource.Sites
+		}
+
+		rCli, err = common.GetResourceClient(rc, ac)
+		if err != nil {
+			return nil, clues.Wrap(err, "creating resource client").WithClues(ctx)
+		}
+	}
+
 	ctrl := Controller{
 		AC:           ac,
 		IDNameLookup: idname.NewCache(nil),
+		ownerLookup:  rCli,
 
 		credentials:        creds,
 		tenant:             acct.ID(),
diff --git a/src/pkg/control/restore_test.go b/src/pkg/control/restore_test.go
index e239c6663..3adc51768 100644
--- a/src/pkg/control/restore_test.go
+++ b/src/pkg/control/restore_test.go
@@ -116,15 +116,18 @@ func (suite *RestoreUnitSuite) TestRestoreConfig_piiHandling() {
 		expectPlain string
 	}{
 		{
-			name:        "empty",
-			expectSafe:  `{"onCollision":"","protectedResource":"","location":"","drive":"","includePermissions":false}`,
-			expectPlain: `{"onCollision":"","protectedResource":"","location":"","drive":"","includePermissions":false}`,
+			name: "empty",
+			expectSafe: `{"onCollision":"","protectedResource":"","SubService":{"ID":"","Type":0}` +
+				`,"location":"","drive":"","includePermissions":false}`,
+			expectPlain: `{"onCollision":"","protectedResource":"","SubService":{"ID":"","Type":0}` +
+				`,"location":"","drive":"","includePermissions":false}`,
 		},
 		{
-			name:       "defaults",
-			rc:         cdrc,
-			expectSafe: `{"onCollision":"skip","protectedResource":"","location":"***","drive":"","includePermissions":false}`,
-			expectPlain: `{"onCollision":"skip","protectedResource":"","location":"` +
+			name: "defaults",
+			rc:   cdrc,
+			expectSafe: `{"onCollision":"skip","protectedResource":"","SubService":{"ID":"","Type":0}` +
+				`,"location":"***","drive":"","includePermissions":false}`,
+			expectPlain: `{"onCollision":"skip","protectedResource":"","SubService":{"ID":"","Type":0},"location":"` +
 				cdrc.Location + `","drive":"","includePermissions":false}`,
 		},
 		{
@@ -136,9 +139,11 @@ func (suite *RestoreUnitSuite) TestRestoreConfig_piiHandling() {
 				Drive:              "somedriveid",
 				IncludePermissions: true,
 			},
-			expectSafe: `{"onCollision":"copy","protectedResource":"***","location":"***/exchange/***/email/***/***/***",` +
+			expectSafe: `{"onCollision":"copy","protectedResource":"***","SubService":{"ID":"","Type":0}` +
+				`,"location":"***/exchange/***/email/***/***/***",` +
 				`"drive":"***","includePermissions":true}`,
-			expectPlain: `{"onCollision":"copy","protectedResource":"snoob","location":"tid/exchange/ro/email/foo/bar/baz",` +
+			expectPlain: `{"onCollision":"copy","protectedResource":"snoob","SubService":{"ID":"","Type":0}` +
+				`,"location":"tid/exchange/ro/email/foo/bar/baz",` +
 				`"drive":"somedriveid","includePermissions":true}`,
 		},
 	}