From 5e1404497f110a3c3dc9a5dd9de45c37c087349c Mon Sep 17 00:00:00 2001
From: amela <amela@ena.one>
Date: Wed, 13 Nov 2024 10:47:49 +0100
Subject: [PATCH] feat(cmd/network): Autodetect chain context if none passed

---
 cmd/network/set_chain_context.go              | 26 ++++++++++++++++---
 docs/network.md                               | 10 +++++--
 ...context.in => 05-set-chain-context-ctx.in} |  0
 ...ntext.out => 05-set-chain-context-ctx.out} |  0
 examples/network/set-chain-context.in.static  |  1 +
 examples/network/set-chain-context.out.static |  0
 6 files changed, 31 insertions(+), 6 deletions(-)
 rename examples/network/{05-set-chain-context.in => 05-set-chain-context-ctx.in} (100%)
 rename examples/network/{05-set-chain-context.out => 05-set-chain-context-ctx.out} (100%)
 create mode 100644 examples/network/set-chain-context.in.static
 create mode 100644 examples/network/set-chain-context.out.static

diff --git a/cmd/network/set_chain_context.go b/cmd/network/set_chain_context.go
index 78582aa7..77cbc499 100644
--- a/cmd/network/set_chain_context.go
+++ b/cmd/network/set_chain_context.go
@@ -1,20 +1,24 @@
 package network
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/spf13/cobra"
 
+	"github.com/oasisprotocol/oasis-sdk/client-sdk/go/config"
+	"github.com/oasisprotocol/oasis-sdk/client-sdk/go/connection"
+
 	cliConfig "github.com/oasisprotocol/cli/config"
 )
 
 var setChainContextCmd = &cobra.Command{
-	Use:   "set-chain-context <name> <chain-context>",
+	Use:   "set-chain-context <name> [chain-context]",
 	Short: "Sets the chain context of the given network",
-	Args:  cobra.ExactArgs(2),
+	Args:  cobra.RangeArgs(1, 2),
 	Run: func(cmd *cobra.Command, args []string) {
 		cfg := cliConfig.Global()
-		name, chainContext := args[0], args[1]
+		name := args[0]
 
 		net := cfg.Networks.All[name]
 		if net == nil {
@@ -22,7 +26,21 @@ var setChainContextCmd = &cobra.Command{
 			return // To make staticcheck happy as it doesn't know CheckErr exits.
 		}
 
-		net.ChainContext = chainContext
+		if len(args) >= 2 {
+			net.ChainContext = args[1]
+		} else {
+			// Connect to the network and query the chain context.
+			network := config.Network{
+				RPC: net.RPC,
+			}
+			ctx := context.Background()
+			conn, err := connection.ConnectNoVerify(ctx, &network)
+			cobra.CheckErr(err)
+			chainCtx, err := conn.Consensus().GetChainContext(ctx)
+			cobra.CheckErr(err)
+			net.ChainContext = chainCtx
+			cobra.CheckErr(net.Validate())
+		}
 
 		err := cfg.Save()
 		cobra.CheckErr(err)
diff --git a/docs/network.md b/docs/network.md
index dd2a4e04..1f2ab9ff 100644
--- a/docs/network.md
+++ b/docs/network.md
@@ -95,7 +95,7 @@ You can also delete network in non-interactive mode format by passing the
 ## Set Network Chain Context {#set-chain-context}
 
 To change the chain context of a network, use
-`network set-chain-context <name> <chain-context>`.
+`network set-chain-context <name> [chain-context]`.
 
 :::caution
 
@@ -110,12 +110,18 @@ documentation.
 
 ![code shell](../examples/network/04-list.out)
 
-![code shell](../examples/network/05-set-chain-context.in)
+![code shell](../examples/network/05-set-chain-context-ctx.in)
 
 ![code shell](../examples/network/06-list.in)
 
 ![code shell](../examples/network/06-list.out)
 
+To automatically detect the chain context, simply omit the `[chain-context]`
+argument. This is especially useful for Localnet, where the chain context
+changes each time you restart the `oasis-net-runner`:
+
+![code shell](../examples/network/set-chain-context.in.static)
+
 [Mainnet]: https://github.com/oasisprotocol/docs/blob/main/docs/node/mainnet/README.md
 [Testnet]: https://github.com/oasisprotocol/docs/blob/main/docs/node/testnet/README.md
 
diff --git a/examples/network/05-set-chain-context.in b/examples/network/05-set-chain-context-ctx.in
similarity index 100%
rename from examples/network/05-set-chain-context.in
rename to examples/network/05-set-chain-context-ctx.in
diff --git a/examples/network/05-set-chain-context.out b/examples/network/05-set-chain-context-ctx.out
similarity index 100%
rename from examples/network/05-set-chain-context.out
rename to examples/network/05-set-chain-context-ctx.out
diff --git a/examples/network/set-chain-context.in.static b/examples/network/set-chain-context.in.static
new file mode 100644
index 00000000..806def8e
--- /dev/null
+++ b/examples/network/set-chain-context.in.static
@@ -0,0 +1 @@
+oasis network set-chain-context mainnet_local
diff --git a/examples/network/set-chain-context.out.static b/examples/network/set-chain-context.out.static
new file mode 100644
index 00000000..e69de29b