From c1499aeb26aab647f4f57ae6ad121ee0a1da3cb6 Mon Sep 17 00:00:00 2001
From: bingyuyap <bingyu.yap.21@gmail.com>
Date: Thu, 6 Jun 2024 00:28:24 +0800
Subject: [PATCH] cloud_functions: convert env var NETWORK to sdk Network

Signed-off-by: bingyuyap <bingyu.yap.21@gmail.com>
---
 cloud_functions/src/computeNTTRateLimits.ts        |  6 +++---
 cloud_functions/src/computeTotalSupplyAndLocked.ts | 14 ++++++--------
 cloud_functions/src/getNTTRateLimits.ts            |  4 ++--
 cloud_functions/src/getTotalSupplyAndLocked.ts     |  7 ++-----
 4 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/cloud_functions/src/computeNTTRateLimits.ts b/cloud_functions/src/computeNTTRateLimits.ts
index 70226674..962ec502 100644
--- a/cloud_functions/src/computeNTTRateLimits.ts
+++ b/cloud_functions/src/computeNTTRateLimits.ts
@@ -1,5 +1,4 @@
 import {
-  assertEnvironmentVariable,
   NTT_MANAGER_CONTRACT,
   NTT_TOKENS,
   NTT_TRANSCEIVER_CONTRACT,
@@ -7,6 +6,7 @@ import {
   getEvmTokenDecimals,
   getSolanaTokenDecimals,
   NTTRateLimit,
+  getNetwork,
 } from '@wormhole-foundation/wormhole-monitor-common';
 import { EvmPlatform, EvmChains } from '@wormhole-foundation/sdk-evm';
 import { SolanaPlatform } from '@wormhole-foundation/sdk-solana';
@@ -17,7 +17,8 @@ import { Storage } from '@google-cloud/storage';
 
 const storage = new Storage();
 let bucketName: string = 'wormhole-ntt-cache';
-if (assertEnvironmentVariable('NETWORK') === 'Testnet') {
+const network = getNetwork();
+if (network === 'Testnet') {
   bucketName = 'wormhole-ntt-cache-testnet';
 }
 
@@ -115,7 +116,6 @@ export async function computeNTTRateLimits(req: any, res: any) {
   }
 
   try {
-    const network = assertEnvironmentVariable('NETWORK') as Network;
     const managerContracts = NTT_MANAGER_CONTRACT[network];
 
     const rateLimits = await Promise.all(
diff --git a/cloud_functions/src/computeTotalSupplyAndLocked.ts b/cloud_functions/src/computeTotalSupplyAndLocked.ts
index e7846a0b..b8a755f9 100644
--- a/cloud_functions/src/computeTotalSupplyAndLocked.ts
+++ b/cloud_functions/src/computeTotalSupplyAndLocked.ts
@@ -1,24 +1,23 @@
 import {
   NTT_MANAGER_CONTRACT,
   NTT_TOKENS,
-  assertEnvironmentVariable,
+  NTTTotalSupplyAndLockedData,
   derivePda,
+  getCustody,
+  getCustodyAmount,
   getEvmTokenDecimals,
   getEvmTotalSupply,
+  getNetwork,
   normalizeToDecimals,
 } from '@wormhole-foundation/wormhole-monitor-common';
 import { PublicKey } from '@solana/web3.js';
-import {
-  getCustody,
-  getCustodyAmount,
-  NTTTotalSupplyAndLockedData,
-} from '@wormhole-foundation/wormhole-monitor-common';
 import { Network, rpc, Chain, chainToChainId } from '@wormhole-foundation/sdk-base';
 import { Storage } from '@google-cloud/storage';
 
 const storage = new Storage();
 let bucketName: string = 'wormhole-ntt-cache';
-if (assertEnvironmentVariable('NETWORK') === 'Testnet') {
+const network = getNetwork();
+if (network === 'Testnet') {
   bucketName = 'wormhole-ntt-cache-testnet';
 }
 
@@ -96,7 +95,6 @@ export async function computeTotalSupplyAndLocked(req: any, res: any) {
   }
 
   try {
-    const network = assertEnvironmentVariable('NETWORK') as Network;
     const totalSupplyAndLocked = await fetchTotalSupplyAndLocked(network);
     await cloudStorageCache.save(JSON.stringify(totalSupplyAndLocked));
     res.status(200).send('Total supply and locked saved');
diff --git a/cloud_functions/src/getNTTRateLimits.ts b/cloud_functions/src/getNTTRateLimits.ts
index 435b76fc..0187b47d 100644
--- a/cloud_functions/src/getNTTRateLimits.ts
+++ b/cloud_functions/src/getNTTRateLimits.ts
@@ -2,7 +2,7 @@ import { Storage } from '@google-cloud/storage';
 import { chainIdToChain } from '@wormhole-foundation/sdk-base';
 import {
   NTTRateLimit,
-  assertEnvironmentVariable,
+  getNetwork,
   normalizeToDecimals,
 } from '@wormhole-foundation/wormhole-monitor-common';
 import { Gauge, register } from 'prom-client';
@@ -33,7 +33,7 @@ async function setCapacityGauge(
 const storage = new Storage();
 
 let bucketName: string = 'wormhole-ntt-cache';
-const network = assertEnvironmentVariable('NETWORK');
+const network = getNetwork();
 if (network === 'Testnet') {
   bucketName = 'wormhole-ntt-cache-testnet';
 }
diff --git a/cloud_functions/src/getTotalSupplyAndLocked.ts b/cloud_functions/src/getTotalSupplyAndLocked.ts
index 3decc9a6..7825af16 100644
--- a/cloud_functions/src/getTotalSupplyAndLocked.ts
+++ b/cloud_functions/src/getTotalSupplyAndLocked.ts
@@ -1,11 +1,8 @@
 import { Storage } from '@google-cloud/storage';
-import {
-  NTTRateLimit,
-  assertEnvironmentVariable,
-} from '@wormhole-foundation/wormhole-monitor-common';
+import { NTTRateLimit, getNetwork } from '@wormhole-foundation/wormhole-monitor-common';
 
 const storage = new Storage();
-const network = assertEnvironmentVariable('NETWORK');
+const network = getNetwork();
 let bucketName: string = 'wormhole-ntt-cache';
 if (network === 'Testnet') {
   bucketName = 'wormhole-ntt-cache-testnet';