From dfe843a824f6ebfac4c5ac84dd204b0badac5e10 Mon Sep 17 00:00:00 2001
From: Philipp Giese <187786+frontendphil@users.noreply.github.com>
Date: Fri, 17 Jan 2025 10:03:45 +0100
Subject: [PATCH] feat(pilot-app): make sure a role update updates also the
 route endpoint (#568)

---
 .../modules/src/updateRolesWaypoint.spec.ts   | 71 +++++++++++++++++++
 packages/modules/src/updateRolesWaypoint.ts   | 36 ++++++++--
 2 files changed, 103 insertions(+), 4 deletions(-)
 create mode 100644 packages/modules/src/updateRolesWaypoint.spec.ts

diff --git a/packages/modules/src/updateRolesWaypoint.spec.ts b/packages/modules/src/updateRolesWaypoint.spec.ts
new file mode 100644
index 00000000..7cd51728
--- /dev/null
+++ b/packages/modules/src/updateRolesWaypoint.spec.ts
@@ -0,0 +1,71 @@
+import { Chain } from '@zodiac/chains'
+import {
+  createMockEndWaypoint,
+  createMockExecutionRoute,
+  createMockOwnsConnection,
+  createMockRoleWaypoint,
+  createMockWaypoints,
+  randomAddress,
+} from '@zodiac/test-utils'
+import { formatPrefixedAddress } from 'ser-kit'
+import { describe, expect, it } from 'vitest'
+import { createEnabledConnection } from './createEnabledConnection'
+import { getWaypoints } from './getWaypoints'
+import { updateRolesWaypoint } from './updateRolesWaypoint'
+
+describe('updateRolesWaypoint', () => {
+  describe('Create', () => {
+    it('changes an existing endpoint to a IS_ENABLED connection', () => {
+      const route = createMockExecutionRoute({
+        waypoints: createMockWaypoints({
+          end: createMockEndWaypoint({
+            connection: createMockOwnsConnection(),
+          }),
+        }),
+      })
+
+      const address = randomAddress()
+
+      const updatedRoute = updateRolesWaypoint(route, {
+        moduleAddress: address,
+        multisend: [],
+        version: 1,
+      })
+
+      const [, endPoint] = getWaypoints(updatedRoute)
+
+      expect(endPoint).toHaveProperty(
+        'connection',
+        createEnabledConnection(formatPrefixedAddress(Chain.ETH, address)),
+      )
+    })
+  })
+
+  describe('Update', () => {
+    it('changes an existing endpoint to a IS_ENABLED connection', () => {
+      const route = createMockExecutionRoute({
+        waypoints: createMockWaypoints({
+          waypoints: [createMockRoleWaypoint()],
+          end: createMockEndWaypoint({
+            connection: createMockOwnsConnection(),
+          }),
+        }),
+      })
+
+      const address = randomAddress()
+
+      const updatedRoute = updateRolesWaypoint(route, {
+        moduleAddress: address,
+        multisend: [],
+        version: 1,
+      })
+
+      const [, endPoint] = getWaypoints(updatedRoute)
+
+      expect(endPoint).toHaveProperty(
+        'connection',
+        createEnabledConnection(formatPrefixedAddress(Chain.ETH, address)),
+      )
+    })
+  })
+})
diff --git a/packages/modules/src/updateRolesWaypoint.ts b/packages/modules/src/updateRolesWaypoint.ts
index c59aead3..18d0d223 100644
--- a/packages/modules/src/updateRolesWaypoint.ts
+++ b/packages/modules/src/updateRolesWaypoint.ts
@@ -1,8 +1,15 @@
 import { invariant } from '@epic-web/invariant'
 import { getChainId } from '@zodiac/chains'
-import type { ExecutionRoute, HexAddress, Waypoints } from '@zodiac/schema'
-import { AccountType } from 'ser-kit'
+import type {
+  ExecutionRoute,
+  HexAddress,
+  Waypoint,
+  Waypoints,
+} from '@zodiac/schema'
+import { AccountType, formatPrefixedAddress } from 'ser-kit'
+import { createEnabledConnection } from './createEnabledConnection'
 import { createRolesWaypoint } from './createRolesWaypoint'
+import { createSafeWaypoint } from './createSafeWaypoint'
 
 type RoleUpdatePayload = {
   moduleAddress: HexAddress
@@ -31,7 +38,7 @@ export const updateRolesWaypoint = (
   if (hasRolesWaypoint(route.waypoints)) {
     const newWaypoints = waypoints.map((waypoint) => {
       if (waypoint.account.type !== AccountType.ROLES) {
-        return waypoint
+        return updateSafeConnection(waypoint, moduleAddress)
       }
 
       return createRolesWaypoint({
@@ -60,10 +67,31 @@ export const updateRolesWaypoint = (
         version,
         from: startingPoint.account.prefixedAddress,
       }),
-      ...waypoints,
+      ...waypoints.map((waypoint) =>
+        updateSafeConnection(waypoint, moduleAddress),
+      ),
     ],
   }
 }
 
 const hasRolesWaypoint = (waypoints: Waypoints) =>
   waypoints.some((waypoint) => waypoint.account.type === AccountType.ROLES)
+
+const updateSafeConnection = (
+  waypoint: Waypoint,
+  moduleAddress: HexAddress,
+): Waypoint => {
+  const { account } = waypoint
+
+  if (account.type !== AccountType.SAFE) {
+    return waypoint
+  }
+
+  return createSafeWaypoint({
+    chainId: account.chain,
+    safe: account.address,
+    connection: createEnabledConnection(
+      formatPrefixedAddress(account.chain, moduleAddress),
+    ),
+  })
+}