scheduler_test.go

package schedulers

import (
	"log"
	"sync"
	"testing"
	"time"
)

type Work struct {
	Sender   int
	Duration time.Duration
	Done     chan<- struct{}
}

type ShareTable struct {
	table map[int]int
	mu    sync.RWMutex
}

func (T *ShareTable) Inc(user int) {
	T.mu.Lock()
	defer T.mu.Unlock()

	if T.table == nil {
		T.table = make(map[int]int)
	}
	T.table[user]++
}

func (T *ShareTable) Get(user int) int {
	T.mu.RLock()
	defer T.mu.RUnlock()

	v, _ := T.table[user]
	return v
}

func testSink(sched *Scheduler, table *ShareTable) {
	sink := sched.NewSink()
	for {
		w := sink.Read()
		switch v := w.(type) {
		case Work:
			// dummy load
			start := time.Now()
			for time.Since(start) < v.Duration {
			}
			table.Inc(v.Sender)
			close(v.Done)
		}
	}
}

func testSource(sched *Scheduler, id int, dur time.Duration) {
	source := sched.NewSource()
	for {
		done := make(chan struct{})
		w := Work{
			Sender:   id,
			Duration: dur,
			Done:     done,
		}
		source.Schedule(w)
		<-done
	}
}

func TestScheduler(t *testing.T) {
	var table ShareTable