From 1399094e43002709a49b047c0fc3f97a3ebaa860 Mon Sep 17 00:00:00 2001
From: Savely Krendelhoff <alicecandyom@proton.me>
Date: Mon, 17 Jul 2023 13:40:30 +0300
Subject: [PATCH] [Chore] Unify interactive queries in tezos-setup

Problem: Some interactive queries use step concept,
while others are done through `yes_or_no`. It's useful
to unify them by using steps only for logging
and for the #568 (non-interactive mode) issue.

Solution: Rewrite delete node data query using steps
instead of `yes_or_no`.
---
 baking/src/tezos_baking/steps.py              | 15 +++++++++++++++
 baking/src/tezos_baking/tezos_setup_wizard.py | 17 ++++++++++++++++-
 baking/src/tezos_baking/wizard_structure.py   |  5 ++---
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/baking/src/tezos_baking/steps.py b/baking/src/tezos_baking/steps.py
index f991dfd14..46c338935 100644
--- a/baking/src/tezos_baking/steps.py
+++ b/baking/src/tezos_baking/steps.py
@@ -215,3 +215,18 @@ def get_ledger_derivation_query(ledgers_derivations, node_endpoint, client_dir):
             ]
         ),
     )
+
+
+replace_key_options = {
+    "no": "Keep the existing key",
+    "yes": "Import a new key and replace the existing one",
+}
+
+replace_key_query = Step(
+    id="replace_key",
+    prompt="Would you like to import a new key and replace this one?",
+    help="It's possible to proceed with the existing baker key, instead of\n"
+    "importing new one.",
+    options=replace_key_options,
+    validator=Validator(validators.enum_range(replace_key_options)),
+)
diff --git a/baking/src/tezos_baking/tezos_setup_wizard.py b/baking/src/tezos_baking/tezos_setup_wizard.py
index 040739e77..49af3e6d4 100644
--- a/baking/src/tezos_baking/tezos_setup_wizard.py
+++ b/baking/src/tezos_baking/tezos_setup_wizard.py
@@ -288,6 +288,20 @@ def get_snapshot_mode_query(config):
     )
 
 
+delete_node_data_options = {
+    "no": "Keep the existing data",
+    "yes": "Remove the data under the tezos node data directory",
+}
+
+delete_node_data_query = Step(
+    id="delete_node_data",
+    prompt="Delete this data and bootstrap the node again?",
+    help="It's possible to proceed with bootstrapping the node using\n"
+    "the existing blockchain data, instead of importing fresh snapshot.",
+    options=delete_node_data_options,
+    validator=Validator(validators.enum_range(delete_node_data_options)),
+)
+
 snapshot_file_query = Step(
     id="snapshot_file",
     prompt="Provide the path to the node snapshot file.",
@@ -401,7 +415,8 @@ def check_blockchain_data(self):
             )
             print("The Tezos node data directory already has some blockchain data:")
             print("\n".join(["- " + os.path.join(node_dir, path) for path in diff]))
-            if yes_or_no("Delete this data and bootstrap the node again? <y/N> ", "no"):
+            self.query_step(delete_node_data_query)
+            if self.config["delete_node_data"] == "yes":
                 # We first stop the node service, because it's possible that it
                 # will re-create some of the files while we go on with the wizard
                 print_and_log("Stopping node service")
diff --git a/baking/src/tezos_baking/wizard_structure.py b/baking/src/tezos_baking/wizard_structure.py
index c44b986c5..0373707af 100644
--- a/baking/src/tezos_baking/wizard_structure.py
+++ b/baking/src/tezos_baking/wizard_structure.py
@@ -230,9 +230,8 @@ def check_baker_account(self):
             )
             print("Its current address is", address)
 
-            return yes_or_no(
-                "Would you like to import a new key and replace this one? <y/N> ", "no"
-            )
+            self.query_step(replace_key_query)
+            return self.config["replace_key"] == "yes"
         else:
             return True