From 2c09868f54baebd6f7bd0c1ed4f889a124fbb2c1 Mon Sep 17 00:00:00 2001
From: Tom Beynon <tom@tombeynon.co.uk>
Date: Tue, 4 Apr 2023 20:39:35 +0100
Subject: [PATCH 1/2] Handle networks not included in networks.json

---
 src/autostake/index.mjs | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/autostake/index.mjs b/src/autostake/index.mjs
index aed494ae..4c3bfdae 100644
--- a/src/autostake/index.mjs
+++ b/src/autostake/index.mjs
@@ -17,6 +17,7 @@ import EthSigner from '../utils/EthSigner.mjs';
 
 export default function Autostake(mnemonic, opts) {
   opts = opts || {}
+  let failed = false
 
   if (!mnemonic) {
     timeStamp('Please provide a MNEMONIC environment variable')
@@ -26,7 +27,10 @@ export default function Autostake(mnemonic, opts) {
   async function run(networkNames, networksOverridePath) {
     const networks = getNetworksData(networksOverridePath)
     for (const name of networkNames) {
-      if (name && !networks.map(el => el.name).includes(name)) return timeStamp('Invalid network name:', name)
+      if (name && !networks.map(el => el.name).includes(name)){
+        // Assume network will be found in Chain Registry
+        networks.push({ name, path: name })
+      }
     }
     const calls = networks.map(data => {
       return async () => {
@@ -37,7 +41,7 @@ export default function Autostake(mnemonic, opts) {
         health.started('⚛')
         const results = await runWithRetry(data, health)
         const { success, skipped } = results[results.length - 1] || {}
-        if(!skipped){
+        if(!skipped && !failed){
           health.log(`Autostake ${success ? 'completed' : 'failed'} after ${results.length} attempt(s)`)
           results.forEach(({networkRunner, error}, index) => {
             health.log(`Attempt ${index + 1}:`)
@@ -76,7 +80,7 @@ export default function Autostake(mnemonic, opts) {
       error = e.message
     }
     runners.push({ success: networkRunner?.didSucceed(), networkRunner, error, failedAddresses })
-    if (!networkRunner?.didSucceed() && !networkRunner?.forceFail && retries < maxRetries) {
+    if (!networkRunner?.didSucceed() && !networkRunner?.forceFail && retries < maxRetries && !failed) {
       await logResults(health, networkRunner, error, `Failed attempt ${retries + 1}/${maxRetries + 1}, retrying in 30 seconds...`)
       await new Promise(r => setTimeout(r, 30 * 1000));
       return await runWithRetry(data, health, retries + 1, runners)
@@ -111,8 +115,12 @@ export default function Autostake(mnemonic, opts) {
     let config = { ...opts }
     try {
       await network.load()
-    } catch {
-      throw new Error('Unable to load network data for', network.name)
+    } catch(e) {
+      if(e.response.status === 404){
+        failed = true
+        throw new Error(`${network.name} not found in Chain Registry`)
+      }
+      throw new Error(`Unable to load network data for ${network.name}`)
     }
 
     timeStamp('Loaded', network.prettyName)

From 98df45f2590401aff8ce0fe32a1661d64f547e01 Mon Sep 17 00:00:00 2001
From: Tom Beynon <tom@tombeynon.co.uk>
Date: Tue, 4 Apr 2023 20:39:51 +0100
Subject: [PATCH 2/2] Skip logs/throttle when grants loaded in a batch

---
 src/autostake/NetworkRunner.mjs | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/autostake/NetworkRunner.mjs b/src/autostake/NetworkRunner.mjs
index d91f6d39..18b4652b 100644
--- a/src/autostake/NetworkRunner.mjs
+++ b/src/autostake/NetworkRunner.mjs
@@ -135,8 +135,10 @@ export default class NetworkRunner {
       }
     })
     let grantedAddresses = await mapSync(grantCalls, this.opts.batchQueries, (batch, index) => {
-      timeStamp('...batch', index + 1)
-      return this.throttleQuery()
+      if(!allGrants){
+        timeStamp('...batch', index + 1)
+        return this.throttleQuery()
+      }
     })
     return _.compact(grantedAddresses.flat())
   }
@@ -314,7 +316,7 @@ export default class NetworkRunner {
   }
 
   async throttleQuery(){
-    if(!this.opts.queryThrottle) return 
+    if(!this.opts.queryThrottle) return
 
     await new Promise(r => setTimeout(r, this.opts.queryThrottle));
   }