github.com/ngicks/gokugen@v0.0.5/middleware/observe/observe_test.go (about)

     1  package observe_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/ngicks/gokugen"
     9  	mock_gokugen "github.com/ngicks/gokugen/__mock"
    10  	"github.com/ngicks/gokugen/middleware/observe"
    11  	"github.com/stretchr/testify/assert"
    12  )
    13  
    14  func sapmleMw(handler gokugen.ScheduleHandlerFn) gokugen.ScheduleHandlerFn {
    15  	return func(ctx gokugen.SchedulerContext) (gokugen.Task, error) {
    16  		return handler(
    17  			gokugen.WrapContext(
    18  				ctx,
    19  				gokugen.WithWorkFnWrapper(func(_ gokugen.SchedulerContext, workFn gokugen.WorkFn) gokugen.WorkFn {
    20  					return func(taskCtx context.Context, scheduled time.Time) (any, error) {
    21  						ret, _ := workFn(taskCtx, scheduled)
    22  						return ret.(string) + "bar", nil
    23  					}
    24  				}),
    25  				gokugen.WithWorkId(":"),
    26  			),
    27  		)
    28  	}
    29  }
    30  
    31  func TestObserve(t *testing.T) {
    32  	_, mockSched, getTrappedTask := mock_gokugen.BuildMockScheduler(t)
    33  
    34  	ma := gokugen.NewMiddlewareApplicator(mockSched)
    35  
    36  	inputCtx := gokugen.NewPlainContext(
    37  		time.Now(),
    38  		func(taskCtx context.Context, scheduled time.Time) (any, error) {
    39  			return "foo", nil
    40  		},
    41  		nil,
    42  	)
    43  
    44  	ma.Use(
    45  		observe.New(
    46  			func(ctx gokugen.SchedulerContext) {
    47  				assert.Equal(t, ctx.(*gokugen.PlainContext), inputCtx)
    48  			},
    49  			func(ret any, err error) {
    50  				assert.Equal(t, ret.(string), "foo")
    51  			},
    52  		).Middleware,
    53  		sapmleMw,
    54  		observe.New(
    55  			func(ctx gokugen.SchedulerContext) {
    56  				workId, _ := gokugen.GetWorkId(ctx)
    57  				assert.Equal(t, workId, ":")
    58  			},
    59  			func(ret any, err error) {
    60  				assert.Equal(t, ret.(string), "foobar")
    61  			},
    62  		).Middleware,
    63  	)
    64  
    65  	ma.Schedule(inputCtx)
    66  	getTrappedTask().Do(context.TODO())
    67  }