diff --git a/CHANGELOG b/CHANGELOG
index 53da0148ff3..b0b80f4abf8 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -25,6 +25,7 @@ v 8.1.0
   - Don't show "Add README" link in an empty repository if user doesn't have access to push (Stan Hu)
   - Fix error preventing displaying of commit data for a directory with a leading dot (Stan Hu)
   - Speed up load times of issue detail pages by roughly 1.5x
+  - Fix CI badge
   - Require CI jobs to be named
   - Fix CI rendering regressions
   - Allow developer to manage builds
diff --git a/app/models/ci/project_status.rb b/app/models/ci/project_status.rb
index b66f1212f23..2d35aeac225 100644
--- a/app/models/ci/project_status.rb
+++ b/app/models/ci/project_status.rb
@@ -27,9 +27,5 @@ module Ci
     def human_status
       status
     end
-
-    def last_commit_for_ref(ref)
-      commits.where(ref: ref).last
-    end
   end
 end
diff --git a/app/services/ci/image_for_build_service.rb b/app/services/ci/image_for_build_service.rb
index b95835ba093..b8d24193035 100644
--- a/app/services/ci/image_for_build_service.rb
+++ b/app/services/ci/image_for_build_service.rb
@@ -1,17 +1,15 @@
 module Ci
   class ImageForBuildService
     def execute(project, params)
-      image_name =
-        if params[:sha]
-          commit = project.commits.find_by(sha: params[:sha])
-          image_for_commit(commit)
-        elsif params[:ref]
-          commit = project.last_commit_for_ref(params[:ref])
-          image_for_commit(commit)
-        else
-          'build-unknown.svg'
+      sha = params[:sha]
+      sha ||=
+        if params[:ref]
+          project.gl_project.commit(params[:ref]).try(:sha)
         end
 
+      commit = project.commits.ordered.find_by(sha: sha)
+      image_name = image_for_commit(commit)
+
       image_path = Rails.root.join('public/ci', image_name)
 
       OpenStruct.new(
diff --git a/spec/services/ci/image_for_build_service_spec.rb b/spec/services/ci/image_for_build_service_spec.rb
index d7242d684c6..cda7d0c4a51 100644
--- a/spec/services/ci/image_for_build_service_spec.rb
+++ b/spec/services/ci/image_for_build_service_spec.rb
@@ -4,8 +4,9 @@ module Ci
   describe ImageForBuildService do
     let(:service) { ImageForBuildService.new }
     let(:project) { FactoryGirl.create(:ci_project) }
-    let(:gl_project) { FactoryGirl.create(:empty_project, gitlab_ci_project: project) }
-    let(:commit) { FactoryGirl.create(:ci_commit, gl_project: gl_project, ref: 'master') }
+    let(:gl_project) { FactoryGirl.create(:project, gitlab_ci_project: project) }
+    let(:commit_sha) { gl_project.commit('master').sha }
+    let(:commit) { gl_project.ensure_ci_commit(commit_sha) }
     let(:build) { FactoryGirl.create(:ci_build, commit: commit) }
 
     describe :execute do