Make TestInit.setup_gitlab_shell and TestInit.setup_gitaly more robust

Signed-off-by: Rémy Coutable <remy@rymai.me>

spec/support/test_env.rb

@@ -3,6 +3,8 @@ require 'rspec/mocks'
 module TestEnv
   extend self
 
+  ComponentFailedToInstallError = Class.new(StandardError)
+
   # When developing the seed repository, comment out the branch you will modify.
   BRANCH_SHA = {
     'signed-commits'                     => '2d1096e',
@@ -127,50 +129,23 @@ module TestEnv
   end
 
   def setup_gitlab_shell
-    puts "\n==> Setting up Gitlab Shell..."
-    start = Time.now
-    gitlab_shell_dir = Gitlab.config.gitlab_shell.path
-    shell_needs_update = component_needs_update?(gitlab_shell_dir,
-      Gitlab::Shell.version_required)
-
-    unless !shell_needs_update || system('rake', 'gitlab:shell:install')
-      puts "\nGitLab Shell failed to install, cleaning up #{gitlab_shell_dir}!\n"
-      FileUtils.rm_rf(gitlab_shell_dir)
-      exit 1
-    end
-
-    puts "    GitLab Shell setup in #{Time.now - start} seconds...\n"
+    component_timed_setup('GitLab Shell',
+      install_dir: Gitlab.config.gitlab_shell.path,
+      version: Gitlab::Shell.version_required,
+      task: 'gitlab:shell:install')
   end
 
   def setup_gitaly
-    puts "\n==> Setting up Gitaly..."
-    start = Time.now
     socket_path = Gitlab::GitalyClient.address('default').sub(/\Aunix:/, '')
     gitaly_dir = File.dirname(socket_path)
 
-    if gitaly_dir_stale?(gitaly_dir)
-      puts "    Gitaly is outdated, cleaning up #{gitaly_dir}!"
-      FileUtils.rm_rf(gitaly_dir)
-    end
+    component_timed_setup('Gitaly',
+      install_dir: gitaly_dir,
+      version: Gitlab::GitalyClient.expected_server_version,
+      task: "gitlab:gitaly:install[#{gitaly_dir}]") do
 
-    gitaly_needs_update = component_needs_update?(gitaly_dir,
-      Gitlab::GitalyClient.expected_server_version)
-
-    unless !gitaly_needs_update || system('rake', "gitlab:gitaly:install[#{gitaly_dir}]")
-      puts "\nGitaly failed to install, cleaning up #{gitaly_dir}!\n"
-      FileUtils.rm_rf(gitaly_dir)
-      exit 1
+      start_gitaly(gitaly_dir)
     end
-
-    start_gitaly(gitaly_dir)
-    puts "    Gitaly setup in #{Time.now - start} seconds...\n"
-  end
-
-  def gitaly_dir_stale?(dir)
-    gitaly_executable = File.join(dir, 'gitaly')
-    return false unless File.exist?(gitaly_executable)
-
-    File.mtime(gitaly_executable) < File.mtime(Rails.root.join('GITALY_SERVER_VERSION'))
   end
 
   def start_gitaly(gitaly_dir)
@@ -325,6 +300,40 @@ module TestEnv
     end
   end
 
+  def component_timed_setup(component, install_dir:, version:, task:)
+    puts "\n==> Setting up #{component}..."
+    start = Time.now
+
+    ensure_component_dir_name_is_correct!(component, install_dir)
+
+    if component_needs_update?(install_dir, version)
+      # Cleanup the component entirely to ensure we start fresh
+      FileUtils.rm_rf(install_dir)
+      unless system('rake', task)
+        raise ComponentFailedToInstallError
+      end
+    end
+
+    yield if block_given?
+
+  rescue ComponentFailedToInstallError
+    puts "\n#{component} failed to install, cleaning up #{install_dir}!\n"
+    FileUtils.rm_rf(install_dir)
+    exit 1
+  ensure
+    puts "    #{component} setup in #{Time.now - start} seconds...\n"
+  end
+
+  def ensure_component_dir_name_is_correct!(component, path)
+    actual_component_dir_name = File.basename(path)
+    expected_component_dir_name = component.parameterize
+
+    unless actual_component_dir_name == expected_component_dir_name
+      puts "    #{component} install dir should be named '#{expected_component_dir_name}', not '#{actual_component_dir_name}' (full install path given was '#{path}')!\n"
+      exit 1
+    end
+  end
+
   def component_needs_update?(component_folder, expected_version)
     version = File.read(File.join(component_folder, 'VERSION')).strip