play-with-docker/scheduler/scheduler_test.go

package scheduler

import (
	"context"
	"io/ioutil"
	"log"
	"os"
	"sync"
	"testing"
	"time"

	"github.com/play-with-docker/play-with-docker/event"
	"github.com/play-with-docker/play-with-docker/pwd"
	"github.com/play-with-docker/play-with-docker/pwd/types"
	"github.com/play-with-docker/play-with-docker/storage"
	"github.com/stretchr/testify/assert"
)

type mockTask struct {
	name string
	run  func(ctx context.Context, instance *types.Instance) error
}

func (m *mockTask) Name() string {
	return m.name
}
func (m *mockTask) Run(ctx context.Context, instance *types.Instance) error {
	return m.run(ctx, instance)
}

func mockStorage() storage.StorageApi {
	tmpfile, err := ioutil.TempFile("", "pwd")
	if err != nil {
		log.Fatal(err)
	}
	tmpfile.Close()
	os.Remove(tmpfile.Name())
	store, _ := storage.NewFileStorage(tmpfile.Name())
	return store
}

func TestNew(t *testing.T) {
	store := mockStorage()

	s := &types.Session{
		Id:        "aaabbbccc",
		ExpiresAt: time.Now().Add(time.Hour),
	}

	i := &types.Instance{
		SessionId: s.Id,
		Name:      "node1",
		IP:        "10.0.0.1",
	}
	err := store.SessionPut(s)
	assert.Nil(t, err)

	err = store.InstancePut(i)
	assert.Nil(t, err)

	sch, err := NewScheduler(store, event.NewLocalBroker(), &pwd.Mock{})
	assert.Nil(t, err)
	assert.Len(t, sch.scheduledSessions, 1)
}

func TestAddTask(t *testing.T) {
	store := mockStorage()
	sch, err := NewScheduler(store, event.NewLocalBroker(), &pwd.Mock{})
	assert.Nil(t, err)

	task := &mockTask{name: "FooBar"}
	err = sch.AddTask(task)
	assert.Nil(t, err)

	err = sch.AddTask(task)
	assert.NotNil(t, err)

	assert.Equal(t, map[string]Task{"FooBar": task}, sch.tasks)
}

func TestRemoveTask(t *testing.T) {
	store := mockStorage()
	sch, err := NewScheduler(store, event.NewLocalBroker(), &pwd.Mock{})
	assert.Nil(t, err)

	task := &mockTask{name: "FooBar"}
	err = sch.AddTask(task)
	assert.Nil(t, err)

	err = sch.RemoveTask(task)
	assert.Nil(t, err)

	err = sch.RemoveTask(task)
	assert.NotNil(t, err)

	assert.Equal(t, map[string]Task{}, sch.tasks)
}

func TestStart(t *testing.T) {
	store := mockStorage()

	s := &types.Session{
		Id:        "aaabbbccc",
		ExpiresAt: time.Now().Add(time.Hour),
	}

	i := &types.Instance{
		SessionId: s.Id,
		Name:      "node1",
		IP:        "10.0.0.1",
	}
	err := store.SessionPut(s)
	assert.Nil(t, err)

	err = store.InstancePut(i)
	assert.Nil(t, err)

	sch, err := NewScheduler(store, event.NewLocalBroker(), &pwd.Mock{})
	assert.Nil(t, err)

	wg := sync.WaitGroup{}
	wg.Add(1)
	ran := false
	task := &mockTask{name: "FooBar", run: func(ctx context.Context, instance *types.Instance) error {
		ran = true
		wg.Done()
		return nil
	}}
	err = sch.AddTask(task)
	assert.Nil(t, err)

	sch.Start()
	defer sch.Stop()
	wg.Wait()
	assert.True(t, ran)
}

func TestScheduleFromEvent(t *testing.T) {
	s := &types.Session{
		Id:        "aaaabbbbcccc",
		ExpiresAt: time.Now().Add(time.Hour),
	}
	lb := event.NewLocalBroker()
	store := mockStorage()
	store.SessionPut(s)
	sch, err := NewScheduler(store, lb, &pwd.Mock{})
	assert.Nil(t, err)

	sch.Start()
	defer sch.Stop()

	lb.Emit(event.SESSION_NEW, s.Id)

	time.Sleep(time.Second)

	assert.Len(t, sch.scheduledSessions, 1)
}

func TestUnscheduleFromEvent(t *testing.T) {
	s := &types.Session{
		Id:        "aaaabbbbcccc",
		ExpiresAt: time.Now().Add(time.Hour),
	}
	lb := event.NewLocalBroker()
	store := mockStorage()
	store.SessionPut(s)
	sch, err := NewScheduler(store, lb, &pwd.Mock{})
	assert.Nil(t, err)

	sch.Start()
	defer sch.Stop()

	lb.Emit(event.SESSION_END, s.Id)

	time.Sleep(time.Second)

	assert.Len(t, sch.scheduledSessions, 0)
}

func TestCloseSession(t *testing.T) {
	_e := event.NewLocalBroker()
	_p := &pwd.Mock{}
	_s := mockStorage()

	s := &types.Session{
		Id:        "aaaabbbbcccc",
		ExpiresAt: time.Now().Add(-1 * time.Hour),
	}
	_p.On("SessionClose", s).Return(nil)

	_s.SessionPut(s)

	sch, err := NewScheduler(_s, _e, _p)
	assert.Nil(t, err)

	sch.Start()
	defer sch.Stop()

	time.Sleep(2 * time.Second)

	_p.AssertExpectations(t)

	_e.Emit(event.SESSION_END, s.Id)
	time.Sleep(time.Second)

	assert.Len(t, sch.scheduledSessions, 0)
}