tng 0.1.6 → 0.2.2

data/bin/tng

@@ -1,15 +1,14 @@
 #!/usr/bin/env ruby
 # frozen_string_literal: true
 
-# Set a clean process name for the terminal
 $0 = "tng"
 
-# Add the lib directory to the load path for development
 lib_path = File.expand_path("../lib", __dir__)
 $LOAD_PATH.unshift(lib_path) unless $LOAD_PATH.include?(lib_path)
 
-# Load the main gem first to ensure the Rust extension is loaded
 require "tng"
+require "tng/services/direct_generation"
+require "tng/services/extract_methods"
 
 require "tty-prompt"
 require "tty-spinner"
@@ -21,6 +20,7 @@ require "tty-option"
 require "pastel"
 require "tng/ui/show_help"
 require "tng/ui/post_install_box"
+require "tng/ui/theme"
 require "tng/ui/system_status_display"
 require "tng/ui/configuration_display"
 require "tng/ui/authentication_warning_display"
@@ -33,39 +33,34 @@ require "tng/api/http_client"
 require "tng/ui/controller_test_flow_display"
 require "tng/ui/model_test_flow_display"
 require "tng/ui/service_test_flow_display"
+require "tng/ui/other_test_flow_display"
 require "tng/analyzers/controller"
 require "tng/analyzers/model"
 require "tng/analyzers/service"
+require "tng/analyzers/other"
 require "tng/services/test_generator"
 require "tng/services/user_app_config"
 require "tng/utils"
 
 class CLI
+  include Tng::Services::ExtractMethods
   include TTY::Option
   include Tng::Utils
 
   usage do
     program "tng"
     desc "LLM-Powered Rails Test Generator"
-    example "Interactive mode (full UI with search & selection):", "bundle exec tng"
-    example "  • Browse and search controllers/models", ""
-    example "  • Select specific files from filtered lists", ""
-    example "  • Choose test generation options", ""
-    example "", ""
-    example "Direct generation (bypass UI):", ""
-    example "Standard format:", "bundle exec tng --type=controller --file=users_controller"
-    example "Short aliases:", "bundle exec tng -t c -f ping"
-    example "Mixed format:", "bundle exec tng -t=c f=ping"
-    example "All equivalent:", "bundle exec tng --type=model --file=user"
-    example "                ", "bundle exec tng -t m -f user"
-    example "                ", "bundle exec tng -t=m f=user"
+    example "Interactive mode (method-specific test generation):", "bundle exec tng"
+    example "  • Browse and search controllers/models/services", ""
+    example "  • Select specific methods from filtered lists", ""
+    example "  • Generate tests for individual methods", ""
   end
 
   option :type do
     short "-t"
     long "--type=TYPE"
-    desc "Test type (controller, model, service)"
-    permit %w[controller model service c m mo s se]
+    desc "Test type (controller, model, service, other)"
+    permit %w[controller model service other c m mo s se o]
   end
 
   option :file do
@@ -100,7 +95,6 @@ class CLI
   end
 
   def start
-    # Preprocess ARGV to handle mixed formats like: -t=value or f=value
     normalized_argv = preprocess_arguments(ARGV.dup)
 
     parse(normalized_argv)
@@ -111,26 +105,20 @@ class CLI
       return
     end
 
-    # Load Rails environment if we're in a Rails application (only for non-help commands)
     rails_loaded = load_rails_environment
 
-    # Initialize config and clients after Rails is loaded
     initialize_config_and_clients if rails_loaded
 
     if params[:type] && params[:file]
-      # Initialize config for direct generation if not already done
-      initialize_config_and_clients unless @config_initialized
       run_direct_generation
     else
-      # Check if we're in a Rails app and if Rails loaded properly
       if find_rails_root && !rails_loaded && !defined?(Rails)
-        puts @pastel.red("❌ Failed to load Rails environment.")
-        puts @pastel.yellow("Please ensure you're running this command from a Rails application directory.")
-        puts @pastel.dim("Try running: bundle exec tng")
+        puts @pastel.decorate("#{Tng::UI::Theme.icon(:error)} Failed to load Rails environment.", Tng::UI::Theme.color(:error))
+        puts @pastel.decorate("Please ensure you're running this command from a Rails application directory.", Tng::UI::Theme.color(:warning))
+        puts @pastel.decorate("Try running: bundle exec tng", Tng::UI::Theme.color(:muted))
         return
       end
 
-      # Check if configuration is set up before proceeding
       return unless check_configuration
 
       check_system_status
@@ -177,21 +165,27 @@ class CLI
   def main_menu
     loop do
       clear_screen
+
       puts DisplayBanner.new(@pastel, Tng::VERSION).render
 
-      choice = @prompt.select(
-        center_text("Select an option:"),
-        cycle: true,
-        per_page: 5,
-        symbols: { marker: "→" }
-      ) do |menu|
-        menu.choice "Generate tests", :tests
-        menu.choice "User stats", :stats
-        menu.choice "About", :about
-        menu.choice "Exit", :exit, key: "q"
+      begin
+        choice = @prompt.select(
+          center_text("Select an option:"),
+          cycle: true,
+          per_page: 5,
+          symbols: { marker: "→" }
+        ) do |menu|
+          menu.choice "Generate tests", :tests
+          menu.choice "User stats", :stats
+          menu.choice "About", :about
+          menu.choice "Exit", :exit, key: "q"
+        end
+
+        handle_choice(choice)
+      rescue TTY::Reader::InputInterrupt, Interrupt
+        choice = :exit
       end
 
-      handle_choice(choice)
       break if choice == :exit
     end
   end
@@ -219,7 +213,8 @@ class CLI
       menu.choice "Controller", :controller
       menu.choice "Model", :model
       menu.choice "Service", :service
-      menu.choice @pastel.cyan("⬅️  Back"), :back
+      menu.choice "Other", :other
+      menu.choice @pastel.decorate("#{Tng::UI::Theme.icon(:back)} Back", Tng::UI::Theme.color(:info)), :back
     end
 
     case choice
@@ -229,24 +224,23 @@ class CLI
       generate_model_tests
     when :service
       generate_service_tests
-    when :route
-      generate_route_tests
+    when :other
+      generate_other_tests
     end
   end
 
   def generate_controller_tests
-    header = @pastel.bright_white.bold("Scanning for controllers...")
+    header = @pastel.decorate("Scanning for controllers...", Tng::UI::Theme.color(:primary))
     puts center_text(header)
 
     spinner = TTY::Spinner.new(
       center_text("[:spinner] Analyzing controllers..."),
       format: :dots,
-      success_mark: @pastel.green("✅"),
-      error_mark: @pastel.red("❌")
+      success_mark: @pastel.decorate(Tng::UI::Theme.icon(:success), Tng::UI::Theme.color(:success)),
+      error_mark: @pastel.decorate(Tng::UI::Theme.icon(:error), Tng::UI::Theme.color(:error))
     )
     spinner.auto_spin
     controllers = Tng::Analyzers::Controller.files_in_dir("app/controllers").map do |file|
-      # Extract the full path relative to app/controllers to build proper namespace
       relative_path = file[:path].gsub(%r{^.*app/controllers/}, "").gsub(".rb", "")
       namespaced_name = relative_path.split("/").map(&:camelize).join("::")
 
@@ -264,7 +258,7 @@ class CLI
       return
     end
 
-    success_msg = @pastel.green("Found #{controllers.length} controllers")
+    success_msg = @pastel.decorate("Found #{controllers.length} controllers", Tng::UI::Theme.color(:success))
     spinner.success(center_text(success_msg))
 
     controller_choice = flow_display.select_controller(controllers)
@@ -274,19 +268,18 @@ class CLI
   end
 
   def generate_model_tests
-    header = @pastel.bright_white.bold("Scanning for models...")
+    header = @pastel.decorate("Scanning for models...", Tng::UI::Theme.color(:primary))
     puts center_text(header)
 
     spinner = TTY::Spinner.new(
       center_text("[:spinner] Analyzing models..."),
       format: :dots,
-      success_mark: @pastel.green("✅"),
-      error_mark: @pastel.red("❌")
+      success_mark: @pastel.decorate(Tng::UI::Theme.icon(:success), Tng::UI::Theme.color(:success)),
+      error_mark: @pastel.decorate(Tng::UI::Theme.icon(:error), Tng::UI::Theme.color(:error))
     )
     spinner.auto_spin
 
     models = Tng::Analyzers::Model.files_in_dir("app/models").map do |file|
-      # Extract the full path relative to app/models to build proper namespace
       relative_path = file[:path].gsub(%r{^.*app/models/}, "").gsub(".rb", "")
       namespaced_name = relative_path.split("/").map(&:camelize).join("::")
 
@@ -304,7 +297,7 @@ class CLI
       return
     end
 
-    success_msg = @pastel.green("Found #{models.length} models")
+    success_msg = @pastel.decorate("Found #{models.length} models", Tng::UI::Theme.color(:success))
     spinner.success(center_text(success_msg))
 
     model_choice = flow_display.select_model(models)
@@ -314,19 +307,18 @@ class CLI
   end
 
   def generate_service_tests
-    header = @pastel.bright_white.bold("Scanning for services...")
+    header = @pastel.decorate("Scanning for services...", Tng::UI::Theme.color(:primary))
     puts center_text(header)
 
     spinner = TTY::Spinner.new(
       center_text("[:spinner] Analyzing services..."),
       format: :dots,
-      success_mark: @pastel.green("✅"),
-      error_mark: @pastel.red("❌")
+      success_mark: @pastel.decorate(Tng::UI::Theme.icon(:success), Tng::UI::Theme.color(:success)),
+      error_mark: @pastel.decorate(Tng::UI::Theme.icon(:error), Tng::UI::Theme.color(:error))
     )
     spinner.auto_spin
 
     services = Tng::Analyzers::Service.files_in_dir.map do |file|
-      # Extract the full path relative to app/services or app/service to build proper namespace
       relative_path = file[:path].gsub(%r{^.*app/services?/}, "").gsub(".rb", "")
       namespaced_name = relative_path.split("/").map(&:camelize).join("::")
 
@@ -344,7 +336,7 @@ class CLI
       return
     end
 
-    success_msg = @pastel.green("Found #{services.length} services")
+    success_msg = @pastel.decorate("Found #{services.length} services", Tng::UI::Theme.color(:success))
     spinner.success(center_text(success_msg))
 
     service_choice = flow_display.select_service(services)
@@ -353,483 +345,384 @@ class CLI
     show_service_test_options(service_choice)
   end
 
-  def generate_test_for_controller(controller)
-    header = @pastel.bright_white.bold("🎯 Generating tests for #{controller[:name]}")
+  def generate_other_tests
+    header = @pastel.decorate("Scanning for other files...", Tng::UI::Theme.color(:primary))
     puts center_text(header)
 
-    # Show controller analysis
-    analysis_msg = @pastel.dim("Analyzing controller structure...")
-    puts center_text(analysis_msg)
-
-    # Center the progress bar
-    progress_padding = [(@terminal_width - 56) / 2, 0].max # Account for progress bar width
-    progress_bar = TTY::ProgressBar.new(
-      "#{" " * progress_padding}[:bar] :percent :current/:total",
-      total: 5,
-      bar_format: :box,
-      complete: @pastel.green("█"),
-      incomplete: @pastel.dim("░"),
-      width: 40
+    spinner = TTY::Spinner.new(
+      center_text("[:spinner] Analyzing other files..."),
+      format: :dots,
+      success_mark: @pastel.decorate(Tng::UI::Theme.icon(:success), Tng::UI::Theme.color(:success)),
+      error_mark: @pastel.decorate(Tng::UI::Theme.icon(:error), Tng::UI::Theme.color(:error))
     )
+    spinner.auto_spin
 
-    step_msg = @pastel.dim("Parsing controller file...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.3)
+    other_files = Tng::Analyzers::Other.files_in_dir.map do |file|
+      relative_path = file[:relative_path].gsub(".rb", "")
+      namespaced_name = relative_path.split("/").map(&:camelize).join("::")
 
-    step_msg = @pastel.dim("Extracting actions...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.3)
+      {
+        name: namespaced_name,
+        path: file[:path],
+        type: file[:type]
+      }
+    end
 
-    step_msg = @pastel.dim("Analyzing dependencies...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.3)
+    flow_display = OtherTestFlowDisplay.new(@prompt, @pastel)
 
-    step_msg = @pastel.dim("Generating test code...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    # test for the controller
-    result = Tng::Services::TestGenerator.new(@http_client).run_for_controller(controller)
-    sleep(0.5)
+    if other_files.empty?
+      spinner.error
+      flow_display.show_no_other_files_message
+      return
+    end
 
-    step_msg = @pastel.dim("Writing test file...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.2)
+    success_msg = @pastel.decorate("Found #{other_files.length} other files", Tng::UI::Theme.color(:success))
+    spinner.success(center_text(success_msg))
 
-    completion_msg = @pastel.green.bold("✅ Test generation completed!")
-    puts center_text(completion_msg)
+    other_choice = flow_display.select_other_file(other_files)
+    return if other_choice == :back
 
-    # Show post-generation menu if test was generated successfully
-    if result && result[:file_path]
-      show_post_generation_menu(result)
-    else
-      @prompt.keypress(center_text(@pastel.dim("Press any key to continue...")))
-    end
+    show_other_test_options(other_choice)
   end
 
   def show_controller_method_selection(controller)
-    header = @pastel.bright_white.bold("🔍 Analyzing methods for #{controller[:name]}")
+    header = @pastel.decorate("#{Tng::UI::Theme.icon(:info)} Analyzing methods for #{controller[:name]}", Tng::UI::Theme.color(:primary))
     puts center_text(header)
 
-    # Extract methods from controller
     methods = extract_controller_methods(controller)
+    flow_display = ControllerTestFlowDisplay.new(@prompt, @pastel)
 
     if methods.empty?
-      flow_display = ControllerTestFlowDisplay.new(@prompt, @pastel)
       flow_display.show_no_methods_message(controller)
       return
     end
 
-    # Show method selection UI
-    flow_display = ControllerTestFlowDisplay.new(@prompt, @pastel)
     method_choice = flow_display.select_controller_method(controller, methods)
 
     return if method_choice == :back
 
-    # Generate test for selected method
     generate_test_for_controller_method(controller, method_choice)
   end
 
-  def extract_controller_methods(controller)
-    methods = Tng::Analyzers::Controller.methods_for_controller(controller[:name])
-    Array(methods)
-  rescue StandardError => e
-    puts center_text(@pastel.red("❌ Error analyzing controller: #{e.message}"))
-    []
-  end
-
   def generate_test_for_controller_method(controller, method_info)
-    header = @pastel.bright_white.bold("🎯 Generating test for #{controller[:name]}##{method_info[:name]}")
+    header = @pastel.decorate(
+      "#{Tng::UI::Theme.icon(:marker)} Generating test for #{controller[:name]}##{method_info[:name]}", Tng::UI::Theme.color(:primary)
+    )
     puts center_text(header)
 
-    # Show method analysis
-    analysis_msg = @pastel.dim("Analyzing method structure...")
+    analysis_msg = @pastel.decorate("Analyzing method structure...", Tng::UI::Theme.color(:muted))
     puts center_text(analysis_msg)
 
-    # Center the progress bar
     progress_padding = [(@terminal_width - 56) / 2, 0].max # Account for progress bar width
     progress_bar = TTY::ProgressBar.new(
       "#{" " * progress_padding}[:bar] :percent :current/:total",
       total: 5,
       bar_format: :box,
-      complete: @pastel.green("█"),
-      incomplete: @pastel.dim("░"),
+      complete: @pastel.decorate("█", Tng::UI::Theme.color(:success)),
+      incomplete: @pastel.decorate("░", Tng::UI::Theme.color(:muted)),
       width: 40
     )
 
-    step_msg = @pastel.dim("Parsing method details...")
+    step_msg = @pastel.decorate("Parsing method details...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Analyzing method dependencies...")
+    step_msg = @pastel.decorate("Analyzing method dependencies...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Extracting method context...")
+    step_msg = @pastel.decorate("Extracting method context...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Generating method test code...")
+    step_msg = @pastel.decorate("Generating method test code...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
-    # Generate test for the specific method
     result = Tng::Services::TestGenerator.new(@http_client).run_for_controller_method(controller, method_info)
     sleep(0.5)
 
-    step_msg = @pastel.dim("Writing test file...")
+    step_msg = @pastel.decorate("Writing test file...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.2)
 
-    completion_msg = @pastel.green.bold("✅ Method test generation completed!")
+    # Show completion message with test count and timing
+    if result && result[:test_count]
+      count_msg = result[:test_count] == 1 ? "1 test" : "#{result[:test_count]} tests"
+      time_msg = result[:generation_time] ? " in #{format_generation_time(result[:generation_time])}" : ""
+      completion_msg = @pastel.decorate(
+        "#{Tng::UI::Theme.icon(:success)} Generated #{count_msg} for controller method#{time_msg}!", Tng::UI::Theme.color(:success)
+      )
+    else
+      completion_msg = @pastel.decorate("#{Tng::UI::Theme.icon(:success)} Method test generation completed!", Tng::UI::Theme.color(:success))
+    end
     puts center_text(completion_msg)
 
-    # Show post-generation menu if test was generated successfully
     if result && result[:file_path]
       show_post_generation_menu(result)
     else
-      @prompt.keypress(center_text(@pastel.dim("Press any key to continue...")))
+      @prompt.keypress(center_text(@pastel.decorate("Press any key to continue...", Tng::UI::Theme.color(:muted))))
     end
   end
 
   def show_controller_test_options(controller_choice)
-    flow_display = ControllerTestFlowDisplay.new(@prompt, @pastel)
-    test_option = flow_display.select_test_option(controller_choice)
-
-    return if test_option == :back
-
-    # Check authentication configuration before proceeding with test generation
     return unless check_authentication_configuration
 
-    case test_option
-    when :all_possible_tests
-      generate_test_for_controller(controller_choice)
-    when :per_method
-      show_controller_method_selection(controller_choice)
-    end
+    show_controller_method_selection(controller_choice)
   end
 
   def show_model_test_options(model_choice)
-    flow_display = ModelTestFlowDisplay.new(@prompt, @pastel)
-    test_option = flow_display.select_test_option(model_choice)
-
-    return if test_option == :back
-
-    # Check authentication configuration before proceeding with test generation
     return unless check_authentication_configuration
 
-    case test_option
-    when :all_possible_tests
-      generate_test_for_model(model_choice)
-    when :per_method
-      show_model_method_selection(model_choice)
-    end
+    show_model_method_selection(model_choice)
   end
 
   def show_service_test_options(service_choice)
-    flow_display = ServiceTestFlowDisplay.new(@prompt, @pastel)
-    test_option = flow_display.select_test_option(service_choice)
-
-    return if test_option == :back
-
-    # Check authentication configuration before proceeding with test generation
     return unless check_authentication_configuration
 
-    case test_option
-    when :all_possible_tests
-      generate_test_for_service(service_choice)
-    when :per_method
-      show_service_method_selection(service_choice)
-    end
-  end
-
-  def generate_test_for_model(model)
-    header = @pastel.bright_white.bold("🎯 Generating tests for #{model[:name]}")
-    puts center_text(header)
-
-    # Show model analysis
-    analysis_msg = @pastel.dim("Analyzing model structure...")
-    puts center_text(analysis_msg)
-
-    # Center the progress bar
-    progress_padding = [(@terminal_width - 56) / 2, 0].max # Account for progress bar width
-    progress_bar = TTY::ProgressBar.new(
-      "#{" " * progress_padding}[:bar] :percent :current/:total",
-      total: 4,
-      bar_format: :box,
-      complete: @pastel.green("█"),
-      incomplete: @pastel.dim("░"),
-      width: 40
-    )
-
-    step_msg = @pastel.dim("Parsing model file...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.3)
-
-    step_msg = @pastel.dim("Analyzing validations...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.3)
-
-    step_msg = @pastel.dim("Analyzing associations...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.3)
-
-    step_msg = @pastel.dim("Generating test file...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.3)
-
-    # test for the model
-    result = Tng::Services::TestGenerator.new(@http_client).run_for_model(model)
-    sleep(0.5)
-
-    step_msg = @pastel.dim("Writing test file...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.2)
-    p result
-    completion_msg = @pastel.green.bold("✅ Test generation completed!")
-    puts center_text(completion_msg)
-
-    # Show post-generation menu if test was generated successfully
-    if result && result[:file_path]
-      show_post_generation_menu(result)
-    else
-      @prompt.keypress(center_text(@pastel.dim("Press any key to continue...")))
-    end
+    show_service_method_selection(service_choice)
   end
 
   def show_model_method_selection(model)
-    header = @pastel.bright_white.bold("🔍 Analyzing methods for #{model[:name]}")
+    header = @pastel.decorate("#{Tng::UI::Theme.icon(:info)} Analyzing methods for #{model[:name]}", Tng::UI::Theme.color(:primary))
     puts center_text(header)
 
-    # Extract methods from model
     methods = extract_model_methods(model)
+    flow_display = ModelTestFlowDisplay.new(@prompt, @pastel)
 
     if methods.empty?
-      flow_display = ModelTestFlowDisplay.new(@prompt, @pastel)
       flow_display.show_no_methods_message(model)
       return
     end
 
-    # Show method selection UI
-    flow_display = ModelTestFlowDisplay.new(@prompt, @pastel)
     method_choice = flow_display.select_model_method(model, methods)
 
     return if method_choice == :back
 
-    # Generate test for selected method
     generate_test_for_model_method(model, method_choice)
   end
 
-  def extract_model_methods(model)
-    methods = Tng::Analyzers::Model.methods_for_model(model[:name])
-    Array(methods)
-  rescue StandardError => e
-    puts center_text(@pastel.red("❌ Error analyzing model: #{e.message}"))
-    []
-  end
-
-  def extract_service_methods(service)
-    methods = Tng::Analyzers::Service.methods_for_service(service[:name])
-    Array(methods)
-  rescue StandardError => e
-    puts center_text(@pastel.red("❌ Error analyzing service: #{e.message}"))
-    []
-  end
-
   def generate_test_for_model_method(model, method_info)
-    header = @pastel.bright_white.bold("🎯 Generating test for #{model[:name]}##{method_info[:name]}")
+    header = @pastel.decorate(
+      "#{Tng::UI::Theme.icon(:marker)} Generating test for #{model[:name]}##{method_info[:name]}", Tng::UI::Theme.color(:primary)
+    )
     puts center_text(header)
 
-    # Show method analysis
-    analysis_msg = @pastel.dim("Analyzing method structure...")
+    analysis_msg = @pastel.decorate("Analyzing method structure...", Tng::UI::Theme.color(:muted))
     puts center_text(analysis_msg)
 
-    # Center the progress bar
     progress_padding = [(@terminal_width - 56) / 2, 0].max # Account for progress bar width
     progress_bar = TTY::ProgressBar.new(
       "#{" " * progress_padding}[:bar] :percent :current/:total",
       total: 4,
       bar_format: :box,
-      complete: @pastel.green("█"),
-      incomplete: @pastel.dim("░"),
+      complete: @pastel.decorate("█", Tng::UI::Theme.color(:success)),
+      incomplete: @pastel.decorate("░", Tng::UI::Theme.color(:muted)),
       width: 40
     )
 
-    step_msg = @pastel.dim("Parsing method details...")
+    step_msg = @pastel.decorate("Parsing method details...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Analyzing method dependencies...")
+    step_msg = @pastel.decorate("Analyzing method dependencies...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Generating method test code...")
+    step_msg = @pastel.decorate("Generating method test code...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
-    # Generate test for the specific method
     result = Tng::Services::TestGenerator.new(@http_client).run_for_model_method(model, method_info)
     sleep(0.5)
 
-    step_msg = @pastel.dim("Writing test file...")
+    step_msg = @pastel.decorate("Writing test file...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.2)
 
-    completion_msg = @pastel.green.bold("✅ Method test generation completed!")
+    # Show completion message with test count and timing
+    if result && result[:test_count]
+      count_msg = result[:test_count] == 1 ? "1 test" : "#{result[:test_count]} tests"
+      time_msg = result[:generation_time] ? " in #{format_generation_time(result[:generation_time])}" : ""
+      completion_msg = @pastel.decorate(
+        "#{Tng::UI::Theme.icon(:success)} Generated #{count_msg} for model method#{time_msg}!", Tng::UI::Theme.color(:success)
+      )
+    else
+      completion_msg = @pastel.decorate("#{Tng::UI::Theme.icon(:success)} Method test generation completed!", Tng::UI::Theme.color(:success))
+    end
     puts center_text(completion_msg)
 
-    # Show post-generation menu if test was generated successfully
     if result && result[:file_path]
       show_post_generation_menu(result)
     else
-      @prompt.keypress(center_text(@pastel.dim("Press any key to continue...")))
+      @prompt.keypress(center_text(@pastel.decorate("Press any key to continue...", Tng::UI::Theme.color(:muted))))
     end
   end
 
   def generate_test_for_service_method(service, method_info)
-    header = @pastel.bright_white.bold("🎯 Generating test for #{service[:name]}##{method_info[:name]}")
+    header = @pastel.decorate(
+      "#{Tng::UI::Theme.icon(:marker)} Generating test for #{service[:name]}##{method_info[:name]}", Tng::UI::Theme.color(:primary)
+    )
     puts center_text(header)
 
-    # Show method analysis
-    analysis_msg = @pastel.dim("Analyzing method structure...")
+    analysis_msg = @pastel.decorate("Analyzing method structure...", Tng::UI::Theme.color(:muted))
     puts center_text(analysis_msg)
 
-    # Center the progress bar
     progress_padding = [(@terminal_width - 56) / 2, 0].max # Account for progress bar width
     progress_bar = TTY::ProgressBar.new(
       "#{" " * progress_padding}[:bar] :percent :current/:total",
       total: 4,
       bar_format: :box,
-      complete: @pastel.green("█"),
-      incomplete: @pastel.dim("░"),
+      complete: @pastel.decorate("█", Tng::UI::Theme.color(:success)),
+      incomplete: @pastel.decorate("░", Tng::UI::Theme.color(:muted)),
       width: 40
     )
 
-    step_msg = @pastel.dim("Parsing method details...")
+    step_msg = @pastel.decorate("Parsing method details...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Analyzing method dependencies...")
+    step_msg = @pastel.decorate("Analyzing method dependencies...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Generating method test code...")
+    step_msg = @pastel.decorate("Generating method test code...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
-    # Generate test for the specific method
     result = Tng::Services::TestGenerator.new(@http_client).run_for_service_method(service, method_info)
     sleep(0.5)
 
-    step_msg = @pastel.dim("Writing test file...")
+    step_msg = @pastel.decorate("Writing test file...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.2)
 
-    completion_msg = @pastel.green.bold("✅ Method test generation completed!")
+    # Show completion message with test count and timing
+    if result && result[:test_count]
+      count_msg = result[:test_count] == 1 ? "1 test" : "#{result[:test_count]} tests"
+      time_msg = result[:generation_time] ? " in #{format_generation_time(result[:generation_time])}" : ""
+      completion_msg = @pastel.decorate(
+        "#{Tng::UI::Theme.icon(:success)} Generated #{count_msg} for service method#{time_msg}!", Tng::UI::Theme.color(:success)
+      )
+    else
+      completion_msg = @pastel.decorate("#{Tng::UI::Theme.icon(:success)} Method test generation completed!", Tng::UI::Theme.color(:success))
+    end
     puts center_text(completion_msg)
 
-    # Show post-generation menu if test was generated successfully
     if result && result[:file_path]
       show_post_generation_menu(result)
     else
-      @prompt.keypress(center_text(@pastel.dim("Press any key to continue...")))
+      @prompt.keypress(center_text(@pastel.decorate("Press any key to continue...", Tng::UI::Theme.color(:muted))))
     end
   end
 
   def show_service_method_selection(service)
-    header = @pastel.bright_white.bold("🔍 Analyzing methods for #{service[:name]}")
+    header = @pastel.decorate("#{Tng::UI::Theme.icon(:info)} Analyzing methods for #{service[:name]}", Tng::UI::Theme.color(:primary))
     puts center_text(header)
 
-    # Extract methods from service
     methods = extract_service_methods(service)
+    flow_display = ServiceTestFlowDisplay.new(@prompt, @pastel)
 
     if methods.empty?
-      flow_display = ServiceTestFlowDisplay.new(@prompt, @pastel)
       flow_display.show_no_methods_message(service)
       return
     end
 
-    # Show method selection UI
-    flow_display = ServiceTestFlowDisplay.new(@prompt, @pastel)
     method_choice = flow_display.select_service_method(service, methods)
 
     return if method_choice == :back
 
-    # Generate test for selected method
     generate_test_for_service_method(service, method_choice)
   end
 
-  def generate_test_for_service(service)
-    header = @pastel.bright_white.bold("🎯 Generating tests for #{service[:name]}")
+  def show_other_test_options(other_choice)
+    return unless check_authentication_configuration
+
+    show_other_method_selection(other_choice)
+  end
+
+  def show_other_method_selection(other_file)
+    header = @pastel.decorate("#{Tng::UI::Theme.icon(:info)} Analyzing methods for #{other_file[:name]}", Tng::UI::Theme.color(:primary))
+    puts center_text(header)
+
+    methods = extract_other_methods(other_file)
+    flow_display = OtherTestFlowDisplay.new(@prompt, @pastel)
+
+    if methods.empty?
+      flow_display.show_no_methods_message(other_file)
+      return
+    end
+
+    method_choice = flow_display.select_other_method(other_file, methods)
+
+    return if method_choice == :back
+
+    generate_test_for_other_method(other_file, method_choice)
+  end
+
+  def generate_test_for_other_method(other_file, method_info)
+    header = @pastel.decorate(
+      "#{Tng::UI::Theme.icon(:marker)} Generating test for #{other_file[:name]}##{method_info[:name]}", Tng::UI::Theme.color(:primary)
+    )
     puts center_text(header)
 
-    # Show service analysis
-    analysis_msg = @pastel.dim("Analyzing service structure...")
+    analysis_msg = @pastel.decorate("Analyzing method structure...", Tng::UI::Theme.color(:muted))
     puts center_text(analysis_msg)
 
-    # Center the progress bar
     progress_padding = [(@terminal_width - 56) / 2, 0].max # Account for progress bar width
     progress_bar = TTY::ProgressBar.new(
       "#{" " * progress_padding}[:bar] :percent :current/:total",
       total: 4,
       bar_format: :box,
-      complete: @pastel.green("█"),
-      incomplete: @pastel.dim("░"),
+      complete: @pastel.decorate("█", Tng::UI::Theme.color(:success)),
+      incomplete: @pastel.decorate("░", Tng::UI::Theme.color(:muted)),
       width: 40
     )
 
-    step_msg = @pastel.dim("Parsing service file...")
+    step_msg = @pastel.decorate("Parsing method details...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Analyzing methods...")
+    step_msg = @pastel.decorate("Analyzing method dependencies...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.3)
 
-    step_msg = @pastel.dim("Analyzing dependencies...")
+    step_msg = @pastel.decorate("Generating method test code...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
-    sleep(0.3)
-
-    step_msg = @pastel.dim("Generating test file...")
-    puts center_text(step_msg)
-    progress_bar.advance(1)
-    sleep(0.3)
-
-    # test for the service
-    result = Tng::Services::TestGenerator.new(@http_client).run_for_service(service)
+    result = Tng::Services::TestGenerator.new(@http_client).run_for_other_method(other_file, method_info)
     sleep(0.5)
 
-    step_msg = @pastel.dim("Writing test file...")
+    step_msg = @pastel.decorate("Writing test file...", Tng::UI::Theme.color(:muted))
     puts center_text(step_msg)
     progress_bar.advance(1)
     sleep(0.2)
 
-    completion_msg = @pastel.green.bold("✅ Test generation completed!")
+    # Show completion message with test count and timing
+    if result && result[:test_count]
+      count_msg = result[:test_count] == 1 ? "1 test" : "#{result[:test_count]} tests"
+      time_msg = result[:generation_time] ? " in #{format_generation_time(result[:generation_time])}" : ""
+      completion_msg = @pastel.decorate(
+        "#{Tng::UI::Theme.icon(:success)} Generated #{count_msg} for other method#{time_msg}!", Tng::UI::Theme.color(:success)
+      )
+    else
+      completion_msg = @pastel.decorate("#{Tng::UI::Theme.icon(:success)} Method test generation completed!", Tng::UI::Theme.color(:success))
+    end
     puts center_text(completion_msg)
 
-    # Show post-generation menu if test was generated successfully
     if result && result[:file_path]
       show_post_generation_menu(result)
     else
-      @prompt.keypress(center_text(@pastel.dim("Press any key to continue...")))
+      @prompt.keypress(center_text(@pastel.decorate("Press any key to continue...", Tng::UI::Theme.color(:muted))))
     end
   end
 
@@ -840,8 +733,10 @@ class CLI
   end
 
   def check_system_status
-    # Only show status check in interactive mode
-    puts @pastel.dim("🔍 Checking system status...") unless params[:type] && params[:file]
+    unless params[:type] && params[:file]
+      puts @pastel.decorate("#{Tng::UI::Theme.icon(:info)} Checking system status...",
+                            Tng::UI::Theme.color(:muted))
+    end
 
     status = @testng.check_system_status
     status_ok = SystemStatusDisplay.new(@pastel, params).display(status)
@@ -849,7 +744,7 @@ class CLI
     return if status_ok
 
     puts
-    puts @pastel.dim("Press any key to exit...")
+    puts @pastel.decorate("Press any key to exit...", Tng::UI::Theme.color(:muted))
     $stdin.getch
     exit(1)
   end
@@ -867,12 +762,8 @@ class CLI
     auth_warning = AuthenticationWarningDisplay.new(@pastel)
     auth_missing = auth_warning.display_if_missing
 
-    if auth_missing
-      # User chose to cancel (pressed Esc)
-      return false
-    end
+    return false if auth_missing
 
-    # User chose to continue (pressed Enter) or no auth issues
     true
   end
 
@@ -888,207 +779,69 @@ class CLI
   end
 
   def run_direct_generation
-    # Normalize type aliases
-    type = normalize_type(params[:type])
-
-    unless %w[controller model service].include?(type)
-      puts @pastel.red("❌ Invalid type: #{params[:type]}")
-      puts @pastel.yellow("Supported types: controller, model, service")
-      return
-    end
-
-    # Check configuration first
     return unless check_configuration
 
-    # Check authentication configuration
     return unless check_authentication_configuration
 
-    # Check system status (but don't show the spinner in direct mode)
     status = @testng.check_system_status
     unless SystemStatusDisplay.new(@pastel, params).display(status)
       puts @pastel.red("❌ System status check failed")
       return
     end
 
-    case type
-    when "controller"
-      run_direct_controller_generation
-    when "model"
-      run_direct_model_generation
-    when "service"
-      run_direct_service_generation
-    end
-  end
-
-  def normalize_type(type_param)
-    case type_param&.downcase
-    when "c", "controller"
-      "controller"
-    when "m", "mo", "model"
-      "model"
-    when "s", "se", "service"
-      "service"
-    else
-      type_param&.downcase
-    end
-  end
-
-  def run_direct_controller_generation
-    file_name = params[:file]
-    file_name += "_controller" unless file_name.end_with?("_controller")
-
-    # Build the expected path
-    controller_path = File.join("app/controllers", "#{file_name}.rb")
-
-    # Check if file exists
-    unless File.exist?(controller_path)
-      puts @pastel.red("❌ Controller file not found: #{controller_path}")
-      return
-    end
-
-    # Extract the full path relative to app/controllers to build proper namespace
-    relative_path = controller_path.gsub(%r{^.*app/controllers/}, "").gsub(".rb", "")
-    namespaced_name = relative_path.split("/").map(&:camelize).join("::")
-
-    controller = {
-      name: namespaced_name,
-      path: controller_path
-    }
-
-    puts @pastel.bright_white("🎯 Generating tests for #{controller[:name]}...")
-
-    # Generate the test directly without progress bars for cleaner output
-    result = Tng::Services::TestGenerator.new(@http_client).run_for_controller(controller)
-
-    if result && result[:file_path]
-      show_post_generation_menu(result)
-    else
-      puts @pastel.red("❌ Failed to generate test")
-    end
-  end
-
-  def run_direct_model_generation
-    file_name = params[:file]
-
-    # Build the expected path
-    model_path = File.join("app/models", "#{file_name}.rb")
-
-    # Check if file exists
-    unless File.exist?(model_path)
-      puts @pastel.red("❌ Model file not found: #{model_path}")
-      return
-    end
-
-    # Extract the full path relative to app/models to build proper namespace
-    relative_path = model_path.gsub(%r{^.*app/models/}, "").gsub(".rb", "")
-    namespaced_name = relative_path.split("/").map(&:camelize).join("::")
-
-    model = {
-      name: namespaced_name,
-      path: model_path
-    }
-
-    puts @pastel.bright_white("🎯 Generating tests for #{model[:name]}...")
-
-    # Generate the test directly without progress bars for cleaner output
-    result = Tng::Services::TestGenerator.new(@http_client).run_for_model(model)
-
-    if result && result[:file_path]
-      show_post_generation_menu(result)
-    else
-      puts @pastel.red("❌ Failed to generate test")
-    end
-  end
-
-  def run_direct_service_generation
-    file_name = params[:file]
-
-    # Try different service directories
-    service_dirs = ["app/services", "app/service"]
-    service_path = nil
-
-    service_dirs.each do |dir|
-      potential_path = File.join(dir, "#{file_name}.rb")
-      if File.exist?(potential_path)
-        service_path = potential_path
-        break
-      end
-    end
-
-    # Check if file exists
-    unless service_path
-      puts @pastel.red("❌ Service file not found: #{file_name}.rb")
-      puts @pastel.yellow("Searched in: #{service_dirs.join(", ")}")
-      return
-    end
-
-    # Extract the full path relative to app/services or app/service to build proper namespace
-    relative_path = service_path.gsub(%r{^.*app/services?/}, "").gsub(".rb", "")
-    namespaced_name = relative_path.split("/").map(&:camelize).join("::")
-
-    service = {
-      name: namespaced_name,
-      path: service_path
-    }
-
-    puts @pastel.bright_white("🎯 Generating tests for #{service[:name]}...")
-
-    # Generate the test directly without progress bars for cleaner output
-    result = Tng::Services::TestGenerator.new(@http_client).run_for_service(service)
+    direct_generator = Tng::Services::DirectGeneration.new(
+      @pastel,
+      @testng,
+      @http_client,
+      params,
+      method(:show_post_generation_menu)
+    )
 
-    if result && result[:file_path]
-      show_post_generation_menu(result)
-    else
-      puts @pastel.red("❌ Failed to generate test")
-    end
+    direct_generator.run
   end
 
   def show_post_generation_menu(result)
-    # Check if we're in direct mode
     is_direct_mode = params[:type] && params[:file]
 
     if is_direct_mode
-      # In direct mode, show info and exit
       puts
-      header = @pastel.bright_white.bold("🎉 Test Generated Successfully!")
+      header = @pastel.decorate("#{Tng::UI::Theme.icon(:rocket)} Test Generated Successfully!", Tng::UI::Theme.color(:primary))
       puts center_text(header)
       puts
 
-      # Show file info
       info_msg = [
-        @pastel.cyan("📁 File: #{result[:file_path]}"),
-        @pastel.yellow("🏃 Run: #{result[:run_command]}")
+        @pastel.decorate("#{Tng::UI::Theme.icon(:config)} File: #{result[:file_path]}", Tng::UI::Theme.color(:info)),
+        @pastel.decorate("#{Tng::UI::Theme.icon(:marker)} Run: #{result[:run_command]}", Tng::UI::Theme.color(:warning))
       ].join("\n")
       puts center_text(info_msg)
       puts
 
-      # In direct mode, just show the success message and exit
-      puts center_text(@pastel.green("✅ Test generation completed successfully!"))
+      puts center_text(@pastel.decorate("#{Tng::UI::Theme.icon(:success)} Test generation completed successfully!", Tng::UI::Theme.color(:success)))
       return
     end
 
-    # Interactive mode - show menu loop
     loop do
       puts
-      header = @pastel.bright_white.bold("🎉 Test Generated Successfully!")
+      header = @pastel.decorate("#{Tng::UI::Theme.icon(:rocket)} Test Generated Successfully!", Tng::UI::Theme.color(:primary))
       puts center_text(header)
       puts
 
-      # Show file info
       info_msg = [
-        @pastel.cyan("📁 File: #{result[:file_path]}"),
-        @pastel.yellow("🏃 Run: #{result[:run_command]}")
+        @pastel.decorate("#{Tng::UI::Theme.icon(:config)} File: #{result[:file_path]}", Tng::UI::Theme.color(:info)),
+        @pastel.decorate("#{Tng::UI::Theme.icon(:marker)} Run: #{result[:run_command]}", Tng::UI::Theme.color(:warning))
       ].join("\n")
       puts center_text(info_msg)
       puts
 
       choice = @prompt.select(
-        @pastel.bright_white("What would you like to do?"),
+        @pastel.decorate("What would you like to do?", Tng::UI::Theme.color(:primary)),
         cycle: true,
-        symbols: { marker: "▶" }
+        symbols: { marker: Tng::UI::Theme.icon(:marker) }
       ) do |menu|
-        menu.choice @pastel.yellow("📋 Copy run command"), :copy_command
-        menu.choice @pastel.cyan("⬅️  Back to main menu"), :back
+        menu.choice @pastel.decorate("#{Tng::UI::Theme.icon(:config)} Copy run command", Tng::UI::Theme.color(:warning)),
+                    :copy_command
+        menu.choice @pastel.decorate("#{Tng::UI::Theme.icon(:back)} Back to main menu", Tng::UI::Theme.color(:info)),
+                    :back
       end
 
       case choice
@@ -1103,7 +856,6 @@ class CLI
   def initialize_config_and_clients
     @config_initialized = true
 
-    # Check if configuration is properly set up
     unless Tng::Services::UserAppConfig.configured?
       missing = Tng::Services::UserAppConfig.missing_config
       ConfigurationDisplay.new(@pastel).display_missing_config(missing)
@@ -1116,6 +868,20 @@ class CLI
     )
     @testng = Services::Testng.new(@http_client)
   end
+
+  private
+
+  def format_generation_time(seconds)
+    if seconds < 1
+      "#{(seconds * 1000).round}ms"
+    elsif seconds < 60
+      "#{seconds.round(1)}s"
+    else
+      minutes = (seconds / 60).floor
+      remaining_seconds = (seconds % 60).round
+      "#{minutes}m #{remaining_seconds}s"
+    end
+  end
 end
 
 cli = CLI.new