From 57cef1eabc1f540a4963020eef14d301aafa8533 Mon Sep 17 00:00:00 2001
From: Kristaps Fabians Geikins <fabians@speckle.systems>
Date: Mon, 29 Jul 2024 15:42:49 +0300
Subject: [PATCH] workspace mocks fix, post server invites + setter helpers
 (#2536)

---
 .../server/modules/shared/helpers/mocks.ts    | 14 ++++++++++++-
 .../graph/directives/hasWorkspaceRole.ts      |  6 +++++-
 .../workspaces/graph/mocks/workspaces.ts      | 20 +++++++++++++++----
 3 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/packages/server/modules/shared/helpers/mocks.ts b/packages/server/modules/shared/helpers/mocks.ts
index aebf4f0ed..dd6d6f870 100644
--- a/packages/server/modules/shared/helpers/mocks.ts
+++ b/packages/server/modules/shared/helpers/mocks.ts
@@ -65,6 +65,14 @@ export const mockStoreHelpers = (store: IMockStore) => {
     return ret as T
   }
 
+  const setMockValues = (ref: Ref | SimpleRef, values: Record<string, any>) => {
+    if (isRef(ref)) {
+      store.set(ref, values)
+    } else {
+      store.set(ref.type, ref.id, values)
+    }
+  }
+
   const getFieldValue = <T = any>(
     refOrObj: Record<string, unknown> | Ref | SimpleRef,
     field: string
@@ -174,7 +182,11 @@ export const mockStoreHelpers = (store: IMockStore) => {
      * Or it may be a MockRef, but for some reason you can't just define the field in the mock definition
      * and need a resolver
      */
-    resolveAndCache
+    resolveAndCache,
+    /**
+     * Set/update values into a mock in the mockstore. Useful in mocked mutations.
+     */
+    setMockValues
   }
 }
 
diff --git a/packages/server/modules/workspaces/graph/directives/hasWorkspaceRole.ts b/packages/server/modules/workspaces/graph/directives/hasWorkspaceRole.ts
index f169b3bba..9a05948e1 100644
--- a/packages/server/modules/workspaces/graph/directives/hasWorkspaceRole.ts
+++ b/packages/server/modules/workspaces/graph/directives/hasWorkspaceRole.ts
@@ -36,7 +36,11 @@ export const hasWorkspaceRole: GraphqlDirectiveBuilder = () => {
             const [parent, , context, info] = args
 
             // Validate stream role only if parent is a Stream type
-            if (['Workspace'].includes(info.parentType?.name) && parent) {
+            if (
+              ['Workspace'].includes(info.parentType?.name) &&
+              parent &&
+              !('$ref' in parent)
+            ) {
               if (!parent.id) {
                 // This should never happen as long as our resolvers always return workspaces with their IDs
                 throw new ForbiddenError(
diff --git a/packages/server/modules/workspaces/graph/mocks/workspaces.ts b/packages/server/modules/workspaces/graph/mocks/workspaces.ts
index 4ce87fd6b..f225a9bf7 100644
--- a/packages/server/modules/workspaces/graph/mocks/workspaces.ts
+++ b/packages/server/modules/workspaces/graph/mocks/workspaces.ts
@@ -19,7 +19,8 @@ const config: SpeckleModuleMocksConfig = FF_WORKSPACES_MODULE_ENABLED
           getMockRef,
           resolveFromMockParent,
           addMockRefValues,
-          resolveAndCache
+          resolveAndCache,
+          setMockValues
         }
       }) => {
         return {
@@ -44,7 +45,15 @@ const config: SpeckleModuleMocksConfig = FF_WORKSPACES_MODULE_ENABLED
                 throw new Error('Fake workspace update error')
               }
 
-              return getMockRef('Workspace', { values: omit(args.input, ['logoUrl']) })
+              setMockValues(
+                {
+                  type: 'Workspace',
+                  id: args.input.id
+                },
+                omit(args.input, ['logoUrl', 'id'])
+              )
+
+              return getMockRef('Workspace', { id: args.input.id })
             },
             updateRole: (_parent, args) => {
               const val = faker.datatype.boolean()
@@ -164,7 +173,8 @@ const config: SpeckleModuleMocksConfig = FF_WORKSPACES_MODULE_ENABLED
             )
           },
           WorkspaceCollaborator: {
-            role: resolveFromMockParent()
+            role: resolveFromMockParent(),
+            user: resolveFromMockParent()
           },
           PendingWorkspaceCollaborator: {
             user: resolveAndCache((parent) => {
@@ -174,7 +184,9 @@ const config: SpeckleModuleMocksConfig = FF_WORKSPACES_MODULE_ENABLED
 
               return getMockRef('LimitedUser', { values: { name: title } })
             }),
-            invitedBy: resolveAndCache(() => getMockRef('LimitedUser'))
+            invitedBy: resolveAndCache(() => getMockRef('LimitedUser')),
+            workspaceName: resolveFromMockParent(),
+            token: resolveFromMockParent()
           },
           Project: {
             workspace: resolveAndCache(() => {