From 204310757f0d809ebb91bf8bb66e231a6f50cfbc Mon Sep 17 00:00:00 2001
From: Christophe <contact@christophedeveaux.com>
Date: Wed, 15 May 2024 15:50:00 +0000
Subject: [PATCH] Fix typing

Allow arbOwnerPublicActions to be either called:
- without parameters (`.extend(arbOwnerPublicActions)`)
- with parameters (`.extend(arbOwnerPublicActions({ arbOsVersion: X}))`)
---
 src/decorators/arbOwnerPublicActions.ts       | 56 +++++++++++++++----
 .../arbOwnerPublicActions.unit.test.ts        | 43 +++++++++++++-
 2 files changed, 85 insertions(+), 14 deletions(-)

diff --git a/src/decorators/arbOwnerPublicActions.ts b/src/decorators/arbOwnerPublicActions.ts
index 0db18430..f7335b5f 100644
--- a/src/decorators/arbOwnerPublicActions.ts
+++ b/src/decorators/arbOwnerPublicActions.ts
@@ -1,4 +1,4 @@
-import { Transport, Chain, PrepareTransactionRequestReturnType, PublicClient } from 'viem';
+import { Transport, Chain, PrepareTransactionRequestReturnType, Client, Account } from 'viem';
 
 import {
   arbOwnerReadContract,
@@ -13,7 +13,7 @@ import {
 import { ArbOSVersions } from '../contracts';
 
 export type ArbOwnerPublicActions<
-  TArbOsVersion extends ArbOSVersions,
+  TArbOsVersion extends ArbOSVersions = 20,
   TChain extends Chain | undefined = Chain | undefined,
 > = {
   arbOwnerReadContract: <TFunctionName extends ArbOwnerPublicFunctionName<TArbOsVersion>>(
@@ -24,26 +24,58 @@ export type ArbOwnerPublicActions<
     TFunctionName extends ArbOwnerPublicFunctionName<TArbOsVersion>,
   >(
     args: ArbOwnerPrepareTransactionRequestParameters<TArbOsVersion, TFunctionName>,
-  ) => Promise<PrepareTransactionRequestReturnType<TChain> & { chainId: number }>;
+  ) => Promise<PrepareTransactionRequestReturnType<TChain>>;
 };
 
+const defaultArbOsVersion = 20;
+
+// arbOsVersion is passed as a parameter `client.extend(arbOwnerPublicActions({ arbOsVersion: 10 }))`
 export function arbOwnerPublicActions<
   TArbOsVersion extends ArbOSVersions,
   TTransport extends Transport = Transport,
   TChain extends Chain | undefined = Chain | undefined,
->({ arbOsVersion }: { arbOsVersion: TArbOsVersion }) {
-  return (
-    client: PublicClient<TTransport, TChain>,
-  ): ArbOwnerPublicActions<TArbOsVersion, TChain> => {
-    return {
-      arbOwnerReadContract: (args) => arbOwnerReadContract(client, { ...args, arbOsVersion }),
-
+  TAccount extends Account | undefined = Account | undefined,
+>(param: {
+  arbOsVersion: TArbOsVersion;
+}): (client: Client) => ArbOwnerPublicActions<TArbOsVersion, TChain>;
+// No parameter are passed `client.extend(arbOwnerPublicActions)`
+export function arbOwnerPublicActions<
+  TArbOsVersion extends ArbOSVersions,
+  TTransport extends Transport = Transport,
+  TChain extends Chain | undefined = Chain | undefined,
+  TAccount extends Account | undefined = Account | undefined,
+>(param: Client<TTransport, TChain, TAccount>): ArbOwnerPublicActions<typeof defaultArbOsVersion, TChain>;
+export function arbOwnerPublicActions<
+  TArbOsVersion extends ArbOSVersions,
+  TTransport extends Transport = Transport,
+  TChain extends Chain | undefined = Chain | undefined,
+  TAccount extends Account | undefined = Account | undefined,
+>(paramOrClient: { arbOsVersion: TArbOsVersion } | Client<TTransport, TChain, TAccount>) {
+  if ('arbOsVersion' in paramOrClient) {
+    const result: (client: Client) => ArbOwnerPublicActions<TArbOsVersion, TChain> = (client) => ({
+      arbOwnerReadContract: (args) =>
+        arbOwnerReadContract(client, { ...args, arbOsVersion: paramOrClient.arbOsVersion }),
       arbOwnerPrepareTransactionRequest: (args) =>
         // @ts-ignore (todo: fix viem type issue)
         arbOwnerPrepareTransactionRequest(client, {
           ...args,
-          arbOsVersion,
+          arbOsVersion: paramOrClient.arbOsVersion,
         }),
-    };
+    });
+
+    return result;
+  }
+
+  const result: ArbOwnerPublicActions<typeof defaultArbOsVersion, TChain> = {
+    arbOwnerReadContract: (args) =>
+      // @ts-ignore (todo: fix viem type issue)
+      arbOwnerReadContract(paramOrClient, { ...args, arbOsVersion: defaultArbOsVersion }),
+    arbOwnerPrepareTransactionRequest: (args) =>
+      // @ts-ignore (todo: fix viem type issue)
+      arbOwnerPrepareTransactionRequest(paramOrClient, {
+        ...args,
+        arbOsVersion: defaultArbOsVersion,
+      }),
   };
+  return result;
 }
diff --git a/src/decorators/arbOwnerPublicActions.unit.test.ts b/src/decorators/arbOwnerPublicActions.unit.test.ts
index 568a5189..05f1f86d 100644
--- a/src/decorators/arbOwnerPublicActions.unit.test.ts
+++ b/src/decorators/arbOwnerPublicActions.unit.test.ts
@@ -5,6 +5,10 @@ import { nitroTestnodeL2 } from '../chains';
 import { arbOwnerPublicActions } from './arbOwnerPublicActions';
 import { generatePrivateKey, privateKeyToAccount } from 'viem/accounts';
 
+const clientWithoutParam = createPublicClient({
+  chain: nitroTestnodeL2,
+  transport: http(),
+}).extend(arbOwnerPublicActions);
 const client10 = createPublicClient({
   chain: nitroTestnodeL2,
   transport: http(),
@@ -19,6 +23,10 @@ const client20 = createPublicClient({
 }).extend(arbOwnerPublicActions({ arbOsVersion: 20 }));
 const randomAccount = privateKeyToAccount(generatePrivateKey());
 const upgradeExecutorAddress = '0x24198F8A339cd3C47AEa3A764A20d2dDaB4D1b5b';
+const client = createPublicClient({
+  chain: nitroTestnodeL2,
+  transport: http(),
+})
 
 describe('Accept function name based on arbOSVersion', async () => {
   it('Version 10', () => {
@@ -86,6 +94,31 @@ describe('Accept function name based on arbOSVersion', async () => {
       }),
     ).rejects.toThrowError(AbiFunctionNotFoundError);
   });
+
+  it('Default version (20)', () => {
+    // arbOwnerPublicActions without params is defaulted to arbOsVersion 20
+    expectTypeOf<
+      typeof clientWithoutParam.arbOwnerReadContract<'getInfraFeeAccount'>
+    >().toBeCallableWith({
+      functionName: 'getInfraFeeAccount',
+    });
+
+    expectTypeOf<
+      typeof clientWithoutParam.arbOwnerPrepareTransactionRequest<'setL1PricingRewardRecipient'>
+    >().toBeCallableWith({
+      functionName: 'setL1PricingRewardRecipient',
+      account: randomAccount.address,
+      upgradeExecutor: upgradeExecutorAddress,
+      args: [randomAccount.address],
+    });
+
+    expect(
+      clientWithoutParam.arbOwnerReadContract({
+        // @ts-expect-error Not available for version 20
+        functionName: 'onlyOnArbOS10',
+      }),
+    ).rejects.toThrowError(AbiFunctionNotFoundError);
+  });
 });
 
 // Those tests won't fail if the return type is wrong
@@ -98,7 +131,6 @@ describe('Type return values for function in multiple versions', () => {
       }),
     ).resolves.toEqualTypeOf<`0x${string}`>();
   });
-
   it('Version 11', () => {
     expectTypeOf(
       client11.arbOwnerReadContract({
@@ -106,11 +138,18 @@ describe('Type return values for function in multiple versions', () => {
       }),
     ).resolves.toEqualTypeOf<bigint>();
   });
-  it('Version 11', () => {
+  it('Version 20', () => {
     expectTypeOf(
       client20.arbOwnerReadContract({
         functionName: 'getAllChainOwners',
       }),
     ).resolves.toEqualTypeOf<readonly `0x${string}`[]>();
   });
+  it('Default version (20)', () => {
+    expectTypeOf(
+      clientWithoutParam.arbOwnerReadContract({
+        functionName: 'getAllChainOwners',
+      }),
+    ).resolves.toEqualTypeOf<readonly `0x${string}`[]>();
+  });
 });