GetForSession now expects a *types.Session instead of a string

scheduler/task/collect_stats_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/play-with-docker/play-with-docker/docker"
 	"github.com/play-with-docker/play-with-docker/event"
 	"github.com/play-with-docker/play-with-docker/pwd/types"
+	"github.com/play-with-docker/play-with-docker/storage"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 )
@@ -19,8 +20,8 @@ type mockSessionProvider struct {
 	mock.Mock
 }
 
-func (m *mockSessionProvider) GetDocker(sessionId string) (docker.DockerApi, error) {
-	args := m.Called(sessionId)
+func (m *mockSessionProvider) GetDocker(session *types.Session) (docker.DockerApi, error) {
+	args := m.Called(session)
 
 	return args.Get(0).(docker.DockerApi), args.Error(1)
 }
@@ -34,8 +35,9 @@ func (nopCloser) Close() error { return nil }
 func TestCollectStats_Name(t *testing.T) {
 	e := &event.Mock{}
 	f := &docker.FactoryMock{}
+	s := &storage.Mock{}
 
-	task := NewCollectStats(e, f)
+	task := NewCollectStats(e, f, s)
 
 	assert.Equal(t, "CollectStats", task.Name())
 	e.M.AssertExpectations(t)
@@ -46,6 +48,7 @@ func TestCollectStats_Run(t *testing.T) {
 	d := &docker.Mock{}
 	e := &event.Mock{}
 	f := &docker.FactoryMock{}
+	s := &storage.Mock{}
 
 	stats := dockerTypes.StatsJSON{}
 	b, _ := json.Marshal(stats)
@@ -53,13 +56,19 @@ func TestCollectStats_Run(t *testing.T) {
 		IP:        "10.0.0.1",
 		Name:      "aaaabbbb_node1",
 		SessionId: "aaaabbbbcccc",
+		Hostname:  "node1",
 	}
 
-	f.On("GetForSession", i.SessionId).Return(d, nil)
+	sess := &types.Session{
+		Id: "aaaabbbbcccc",
+	}
+
+	s.On("SessionGet", i.SessionId).Return(sess, nil)
+	f.On("GetForSession", sess).Return(d, nil)
 	d.On("GetContainerStats", i.Name).Return(nopCloser{bytes.NewReader(b)}, nil)
 	e.M.On("Emit", CollectStatsEvent, "aaaabbbbcccc", []interface{}{InstanceStats{Instance: i.Name, Mem: "0.00% (0B / 0B)", Cpu: "0.00%"}}).Return()
 
-	task := NewCollectStats(e, f)
+	task := NewCollectStats(e, f, s)
 	ctx := context.Background()
 
 	err := task.Run(ctx, i)