diff --git a/mocks/Datastore.go b/datastore/mock/datastore.go
similarity index 90%
rename from mocks/Datastore.go
rename to datastore/mock/datastore.go
index 4d05f93d..36e74c9b 100644
--- a/mocks/Datastore.go
+++ b/datastore/mock/datastore.go
@@ -1,14 +1,20 @@
 package mocks
 
-import "github.com/stretchr/testify/mock"
+import (
+	"io"
 
-import "io"
-import "github.com/drone/drone/common"
+	"github.com/drone/drone/common"
+	"github.com/stretchr/testify/mock"
+)
 
 type Datastore struct {
 	mock.Mock
 }
 
+func New() *Datastore {
+	return new(Datastore)
+}
+
 func (m *Datastore) User(_a0 string) (*common.User, error) {
 	ret := m.Called(_a0)
 
@@ -230,6 +236,17 @@ func (m *Datastore) BuildLast(_a0 string) (*common.Build, error) {
 
 	return r0, r1
 }
+func (m *Datastore) BuildAgent(_a0 string, _a1 int) (*common.Agent, error) {
+	ret := m.Called(_a0, _a1)
+
+	var r0 *common.Agent
+	if ret.Get(0) != nil {
+		r0 = ret.Get(0).(*common.Agent)
+	}
+	r1 := ret.Error(1)
+
+	return r0, r1
+}
 func (m *Datastore) SetBuild(_a0 string, _a1 *common.Build) error {
 	ret := m.Called(_a0, _a1)
 
@@ -302,3 +319,11 @@ func (m *Datastore) SetBuildTask(_a0 string, _a1 int, _a2 *common.Task) error {
 
 	return r0
 }
+
+func (m *Datastore) SetBuildAgent(_a0 string, _a1 int, _a2 *common.Agent) error {
+	return m.Called(_a0, _a1, _a2).Error(0)
+}
+
+func (m *Datastore) DelBuildAgent(_a0 string, _a1 int) error {
+	return m.Called(_a0, _a1).Error(0)
+}
diff --git a/server/badge_test.go b/server/badge_test.go
index b498218a..1452c043 100644
--- a/server/badge_test.go
+++ b/server/badge_test.go
@@ -11,7 +11,8 @@ import (
 	"github.com/drone/drone/common"
 	"github.com/drone/drone/common/ccmenu"
 	"github.com/drone/drone/datastore"
-	"github.com/drone/drone/mocks"
+	"github.com/drone/drone/datastore/mock"
+	"github.com/drone/drone/server/recorder"
 	. "github.com/franela/goblin"
 	"github.com/gin-gonic/gin"
 )
@@ -34,9 +35,9 @@ func TestBadge(t *testing.T) {
 		g.AfterEach(func() {
 		})
 
-		cycleStateTester := func(expector gin.HandlerFunc, handle gin.HandlerFunc, validator func(state string, w *ResponseRecorder)) {
+		cycleStateTester := func(expector gin.HandlerFunc, handle gin.HandlerFunc, validator func(state string, w *recorder.ResponseRecorder)) {
 			for idx, state := range []string{"", common.StateError, common.StateFailure, common.StateKilled, common.StatePending, common.StateRunning, common.StateSuccess} {
-				w := NewResponseRecorder()
+				w := recorder.NewResponseRecorder()
 				ctx.Writer = w
 
 				repo.Last = &common.Build{
@@ -60,7 +61,7 @@ func TestBadge(t *testing.T) {
 		g.It("should provide SVG response", func() {
 			{
 				// 1. verify no "last" build
-				w := NewResponseRecorder()
+				w := recorder.NewResponseRecorder()
 				ctx.Writer = w
 				ctx.Request.URL.Path += "/status.svg"
 
@@ -72,7 +73,7 @@ func TestBadge(t *testing.T) {
 			}
 
 			// 2. verify a variety of "last" build states
-			cycleStateTester(nil, GetBadge, func(state string, w *ResponseRecorder) {
+			cycleStateTester(nil, GetBadge, func(state string, w *recorder.ResponseRecorder) {
 				g.Assert(w.Status()).Equal(200)
 				g.Assert(w.HeaderMap.Get("content-type")).Equal("image/svg+xml")
 
@@ -95,7 +96,7 @@ func TestBadge(t *testing.T) {
 		g.It("should provide CCTray response", func() {
 			{
 				// 1. verify no "last" build
-				w := NewResponseRecorder()
+				w := recorder.NewResponseRecorder()
 				ctx.Writer = w
 				ctx.Request.URL.Path += "/cc.xml"
 
@@ -116,7 +117,7 @@ func TestBadge(t *testing.T) {
 				ds.On("BuildLast", fullName).Return(repo.Last, nil).Once()
 			},
 				GetCC,
-				func(state string, w *ResponseRecorder) {
+				func(state string, w *recorder.ResponseRecorder) {
 					g.Assert(w.Status()).Equal(200)
 
 					v := ccmenu.CCProjects{}
diff --git a/server/responserecorder_test.go b/server/recorder/recorder.go
similarity index 97%
rename from server/responserecorder_test.go
rename to server/recorder/recorder.go
index 1ab5e1a9..22da43cd 100644
--- a/server/responserecorder_test.go
+++ b/server/recorder/recorder.go
@@ -1,4 +1,4 @@
-package server
+package recorder
 
 import (
 	"bufio"