From 167a0c7f0bd88e565f580e144b59d9eebe7e24f6 Mon Sep 17 00:00:00 2001
From: Douglas Barbosa Alexandre <dbalexandre@gmail.com>
Date: Thu, 4 Aug 2016 17:15:00 -0300
Subject: Remove SHA suffix for removed branches name when importing PR from GH

---
 lib/gitlab/github_import/branch_formatter.rb       |  2 +-
 lib/gitlab/github_import/importer.rb               | 54 +++++++++++-----------
 .../gitlab/github_import/branch_formatter_spec.rb  |  8 +---
 3 files changed, 28 insertions(+), 36 deletions(-)

diff --git a/lib/gitlab/github_import/branch_formatter.rb b/lib/gitlab/github_import/branch_formatter.rb
index 7d2d545b84e..4be4fc8fe37 100644
--- a/lib/gitlab/github_import/branch_formatter.rb
+++ b/lib/gitlab/github_import/branch_formatter.rb
@@ -8,7 +8,7 @@ module Gitlab
       end
 
       def name
-        @name ||= exists? ? ref : "#{ref}-#{short_id}"
+        ref
       end
 
       def valid?
diff --git a/lib/gitlab/github_import/importer.rb b/lib/gitlab/github_import/importer.rb
index 294e44b7124..9ddc8905bd6 100644
--- a/lib/gitlab/github_import/importer.rb
+++ b/lib/gitlab/github_import/importer.rb
@@ -12,7 +12,6 @@ module Gitlab
 
         if credentials
           @client = Client.new(credentials[:user])
-          @formatter = Gitlab::ImportFormatter.new
         else
           raise Projects::ImportService::Error, "Unable to find project import data credentials for project ID: #{@project.id}"
         end
@@ -69,43 +68,42 @@ module Gitlab
         pull_requests = client.pull_requests(repo, state: :all, sort: :created, direction: :asc, per_page: 100)
         pull_requests = pull_requests.map { |raw| PullRequestFormatter.new(project, raw) }.select(&:valid?)
 
-        source_branches_removed = pull_requests.reject(&:source_branch_exists?).map { |pr| [pr.source_branch_name, pr.number] }
-        target_branches_removed = pull_requests.reject(&:target_branch_exists?).map { |pr| [pr.target_branch_name, pr.target_branch_sha] }
-        branches_removed = source_branches_removed | target_branches_removed
-
-        restore_source_branches(source_branches_removed)
-        restore_target_branches(target_branches_removed)
-
         pull_requests.each do |pull_request|
-          merge_request = pull_request.create!
-          apply_labels(merge_request)
-          import_comments(merge_request)
-          import_comments_on_diff(merge_request)
+          begin
+            restore_source_branch(pull_request) unless pull_request.source_branch_exists?
+            restore_target_branch(pull_request) unless pull_request.target_branch_exists?
+
+            merge_request = pull_request.create!
+            apply_labels(merge_request)
+            import_comments(merge_request)
+            import_comments_on_diff(merge_request)
+          rescue ActiveRecord::RecordInvalid => e
+            raise Projects::ImportService::Error, e.message
+          ensure
+            clean_up_restored_branches(pull_request)
+          end
         end
 
         true
-      rescue ActiveRecord::RecordInvalid => e
-        raise Projects::ImportService::Error, e.message
-      ensure
-        clean_up_restored_branches(branches_removed)
       end
 
-      def restore_source_branches(branches)
-        branches.each do |name, number|
-          project.repository.fetch_ref(repo_url, "pull/#{number}/head", name)
-        end
+      def restore_source_branch(pull_request)
+        project.repository.fetch_ref(repo_url, "pull/#{pull_request.number}/head", pull_request.source_branch_name)
       end
 
-      def restore_target_branches(branches)
-        branches.each do |name, sha|
-          project.repository.create_branch(name, sha)
-        end
+      def restore_target_branch(pull_request)
+        project.repository.create_branch(pull_request.target_branch_name, pull_request.target_branch_sha)
       end
 
-      def clean_up_restored_branches(branches)
-        branches.each do |name, _|
-          project.repository.delete_branch(name) rescue Rugged::ReferenceError
-        end
+      def remove_branch(name)
+        project.repository.delete_branch(name)
+      rescue Rugged::ReferenceError
+        nil
+      end
+
+      def clean_up_restored_branches(pull_request)
+        remove_branch(pull_request.source_branch_name) unless pull_request.source_branch_exists?
+        remove_branch(pull_request.target_branch_name) unless pull_request.target_branch_exists?
 
         project.repository.after_remove_branch
       end
diff --git a/spec/lib/gitlab/github_import/branch_formatter_spec.rb b/spec/lib/gitlab/github_import/branch_formatter_spec.rb
index fc9d5204148..e01a80054a3 100644
--- a/spec/lib/gitlab/github_import/branch_formatter_spec.rb
+++ b/spec/lib/gitlab/github_import/branch_formatter_spec.rb
@@ -33,17 +33,11 @@ describe Gitlab::GithubImport::BranchFormatter, lib: true do
   end
 
   describe '#name' do
-    it 'returns raw ref when branch exists' do
+    it 'returns raw ref' do
       branch = described_class.new(project, double(raw))
 
       expect(branch.name).to eq 'feature'
     end
-
-    it 'returns formatted ref when branch does not exist' do
-      branch = described_class.new(project, double(raw.merge(ref: 'removed-branch', sha: '2e5d3239642f9161dcbbc4b70a211a68e5e45e2b')))
-
-      expect(branch.name).to eq 'removed-branch-2e5d3239'
-    end
   end
 
   describe '#repo' do
-- 
cgit v1.2.1