simplify upgrade code.

1 files changed, 13 insertions, 20 deletions

diff --git a/lib/tfa/cli.rb b/lib/tfa/cli.rb
index 159ad73..af10f66 100644
--- a/lib/tfa/cli.rb
+++ b/lib/tfa/cli.rb
@@ -35,35 +35,20 @@ module TFA
 
     desc "upgrade", "upgrade the pstore database to a yml database."
     def upgrade
-      pstore_path = File.join(directory, ".#{filename}.pstore")
-      yml_path = File.join(directory, ".#{filename}.yml")
-
       if !File.exist?(pstore_path)
-        say "Unable to detect #{pstore_path}"
+        say_status :error, "Unable to detect #{pstore_path}"
         return ""
       end
 
-      say "Detected #{pstore_path}"
-      case ask "Would you like to upgrade to #{yml_path}", limited_to: ["yes", "no"]
-      when "yes"
-        say "Let's begin..."
-        pstore_storage = Storage.new(pstore_path)
-        yaml_storage = Storage.new(yml_path)
+      if yes? "Upgrade to #{yml_path}?"
         pstore_storage.each do |row|
           row.each do |name, secret|
-            case ask "Would you like to migrate `#{name}`?", limited_to: ["yes", "no"]
-            when "yes"
-              say "Migrating `#{name}`..."
-              yaml_storage.save(name, secret)
-            end
+            yaml_storage.save(name, secret) if yes?("Migrate `#{name}`?")
           end
         end
-        case ask "Would you like to delete `#{pstore_path}`? (this action cannot be undone.)", limited_to: ["yes", "no"]
-        when "yes"
+        if yes? "Delete `#{pstore_path}`?"
           File.delete(pstore_path)
         end
-      else
-        say "Nothing to do. Goodbye!"
       end
       ""
     end
@@ -71,7 +56,15 @@ module TFA
     private
 
     def storage
-      @storage ||= Storage.new(File.exist?(pstore_path) ? pstore_path : yml_path)
+      File.exist?(pstore_path) ? pstore_storage : yaml_storage
+    end
+
+    def pstore_storage
+      @pstore_storage ||= Storage.new(pstore_path)
+    end
+
+    def yaml_storage
+      @yaml_storage ||= Storage.new(yml_path)
     end
 
     def filename