worker_test.go

package gocelery

import (
    "math/rand"
    "testing"
    "time"
)

// add is test task method
func add(a int, b int) int {
    return a + b
}

// newCeleryWorker creates celery worker
func newCeleryWorker(numWorkers int) *CeleryWorker {
    broker := NewRedisCeleryBroker("localhost", 6379, 0, "")
    backend := NewRedisCeleryBackend("localhost", 6379, 0, "")
    celeryWorker := NewCeleryWorker(broker, backend, numWorkers)
    return celeryWorker
}

// registerTask registers add test task
func registerTask(celeryWorker *CeleryWorker) string {
    taskName := "add"
    registeredTask := add
    celeryWorker.Register(taskName, registeredTask)
    return taskName
}

func TestRegisterTask(t *testing.T) {
    celeryWorker := newCeleryWorker(1)
    taskName := registerTask(celeryWorker)
    receivedTask := celeryWorker.GetTask(taskName)
    if receivedTask == nil {
        t.Errorf("failed to retrieve task")
    }
}

func TestRunTask(t *testing.T) {
    celeryWorker := newCeleryWorker(1)
    taskName := registerTask(celeryWorker)

    // prepare args
    args := []interface{}{
        rand.Int(),
        rand.Int(),
    }

    // Run task normally
    res := add(args[0].(int), args[1].(int))

    // construct task message
    taskMessage := &CeleryTask{
        Id:      generateUUID(),
        Task:    taskName,
        Args:    args,
        Kwargs:  nil,
        Retries: 1,
        ETA:     time.Now().Add(100 * time.Second),
    }
    resultMsg, err := celeryWorker.RunTask(taskMessage)
    if err != nil {
        t.Errorf("failed to run celery task %v: %v", taskMessage, err)
    }

    reflectRes := resultMsg.Result.(int64)

    // check result
    if int64(res) != reflectRes {
        t.Errorf("reflect result %v is different from normal result %v", reflectRes, res)
    }
}

func TestNumWorkers(t *testing.T) {
    numWorkers := rand.Intn(10)
    celeryWorker := newCeleryWorker(numWorkers)
    celeryNumWorkers := celeryWorker.GetNumWorkers()
    if numWorkers != celeryNumWorkers {
        t.Errorf("number of workers are different: %d vs %d", numWorkers, celeryNumWorkers)
    }
}

func TestStartStop(t *testing.T) {
    numWorkers := rand.Intn(10)
    celeryWorker := newCeleryWorker(numWorkers)
    _ = registerTask(celeryWorker)
    go celeryWorker.StartWorker()
    time.Sleep(100 * time.Millisecond)
    celeryWorker.StopWorker()
}