From ea58ea196601193938b542befb50ba0624730c0f Mon Sep 17 00:00:00 2001
From: Chase Fleming <1666730+chasefleming@users.noreply.github.com>
Date: Mon, 26 Feb 2024 15:18:59 -0800
Subject: [PATCH] Add address type check for arg

---
 components/FundAccountSubmitted.tsx | 10 +++++++++-
 lib/common.ts                       |  8 ++++++++
 lib/flow/fund.ts                    | 10 +---------
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/components/FundAccountSubmitted.tsx b/components/FundAccountSubmitted.tsx
index 91a5269..c6d38b7 100644
--- a/components/FundAccountSubmitted.tsx
+++ b/components/FundAccountSubmitted.tsx
@@ -8,6 +8,7 @@ import {useEffect, useState} from "react"
 import {sendScript} from "../lib/flow/send"
 import fcl from "@onflow/fcl"
 import t from "@onflow/types"
+import {getAddressType} from "../lib/common"
 
 const styles: Record<string, ThemeUICSSObject> = {
   resultsContainer: {
@@ -70,9 +71,16 @@ access(all) fun main(account: Address): UFix64 {
     const fetchBalance = async (addr: string) => {
       try {
         setIsFetchingBalance(true)
+
+        const addressArgType =
+          publicConfig.network === "previewnet" &&
+          getAddressType(result.address) === "FLOW"
+            ? t.Address
+            : t.String
+
         const balance = await sendScript({
           script: balanceScript,
-          args: [fcl.arg(addr, t.Address)],
+          args: [fcl.arg(addr, addressArgType)],
         })
         setBalance(balance)
       } catch (error) {
diff --git a/lib/common.ts b/lib/common.ts
index 2d605d9..8ff9d4e 100644
--- a/lib/common.ts
+++ b/lib/common.ts
@@ -16,3 +16,11 @@ export function verifyAPIKey(req: string, keys: string[]): boolean {
   }
   return false
 }
+
+export const getAddressType = function (address: string): "FLOW" | "FLOWEVM" {
+  if (address.length <= 18) {
+    return "FLOW"
+  } else {
+    return "FLOWEVM"
+  }
+}
diff --git a/lib/flow/fund.ts b/lib/flow/fund.ts
index 0a94214..d93d58c 100644
--- a/lib/flow/fund.ts
+++ b/lib/flow/fund.ts
@@ -3,6 +3,7 @@ import * as t from "@onflow/types"
 import {FLOW_TYPE} from "../constants"
 import publicConfig, {TOKEN_FUNDING_AMOUNTS} from "../publicConfig"
 import {sendTransaction} from "./send"
+import {getAddressType} from "../common"
 
 const txFundAccountFLOW = `
 import FlowToken from ${publicConfig.contractFlowToken}
@@ -86,15 +87,6 @@ export const tokens: Tokens = {
   FLOW: {tx: txFundAccountFLOW, amount: TOKEN_FUNDING_AMOUNTS[FLOW_TYPE]},
   FLOWEVM: {tx: txFundAccountFlowEVM, amount: TOKEN_FUNDING_AMOUNTS[FLOW_TYPE]},
 }
-
-function getAddressType(address: string): "FLOW" | "FLOWEVM" {
-  if (address.length <= 18) {
-    return "FLOW"
-  } else {
-    return "FLOWEVM"
-  }
-}
-
 export async function fundAccount(
   address: string,
   token: TokenType,