github.com/machinefi/w3bstream@v1.6.5-rc9.0.20240426031326-b8c7c4876e72/pkg/depends/kit/mq/z_task_test.go (about)

     1  package mq_test
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"os"
     8  	"testing"
     9  	"time"
    10  
    11  	. "github.com/onsi/gomega"
    12  
    13  	"github.com/machinefi/w3bstream/pkg/depends/kit/kit"
    14  	"github.com/machinefi/w3bstream/pkg/depends/kit/metax"
    15  	"github.com/machinefi/w3bstream/pkg/depends/kit/mq"
    16  	"github.com/machinefi/w3bstream/pkg/depends/kit/mq/mem_mq"
    17  )
    18  
    19  func NewTask(subject, id string, arg ...interface{}) *Task {
    20  	t := &Task{}
    21  	t.SetSubject(subject)
    22  
    23  	if id != "" {
    24  		t.SetID(id)
    25  	}
    26  
    27  	if len(arg) > 0 {
    28  		t.SetArg(arg[0])
    29  	}
    30  	return t
    31  }
    32  
    33  type Task struct {
    34  	mq.TaskHeader
    35  	argv string
    36  }
    37  
    38  func (t *Task) Arg() interface{} { return t.argv }
    39  
    40  func (t *Task) SetArg(v interface{}) {
    41  	switch argv := v.(type) {
    42  	case string:
    43  		t.argv = argv
    44  	case []byte:
    45  		t.argv = string(argv)
    46  	}
    47  }
    48  
    49  var (
    50  	managers = []mq.TaskManager{mem_mq.New(10000) /*redis_mq.New()*/}
    51  	channel  = "cc"
    52  	router   = kit.NewRouter()
    53  )
    54  
    55  func init() {
    56  	for _, m := range managers {
    57  		_ = m.Clear(channel)
    58  	}
    59  }
    60  
    61  type OpA struct{}
    62  
    63  func (a *OpA) Output(ctx context.Context) (interface{}, error) {
    64  	fmt.Println(metax.GetMetaFrom(ctx))
    65  	return nil, nil
    66  }
    67  
    68  type OpB struct {
    69  	bytes.Buffer
    70  }
    71  
    72  func (b *OpB) SetArg(v interface{}) error {
    73  	var data []byte
    74  
    75  	switch arg := v.(type) {
    76  	case []byte:
    77  		data = arg
    78  	case string:
    79  		data = []byte(arg)
    80  	default:
    81  		return nil
    82  	}
    83  	_, err := b.Write(data)
    84  	return err
    85  }
    86  
    87  func (b *OpB) Output(ctx context.Context) (interface{}, error) {
    88  	fmt.Println(metax.GetMetaFrom(ctx), b.String())
    89  	return nil, nil
    90  }
    91  
    92  func TestTaskWorker(t *testing.T) {
    93  	for _, tm := range managers {
    94  		tb := mq.NewTaskBoard(tm)
    95  		n := 100
    96  
    97  		for i := 0; i < n; i++ {
    98  			for j := 0; j < 5; j++ {
    99  				_ = tb.Dispatch(
   100  					channel,
   101  					NewTask("OpA", fmt.Sprintf("OpA%d", i), []byte("A")),
   102  				)
   103  				_ = tb.Dispatch(
   104  					channel,
   105  					NewTask("OpB", fmt.Sprintf("OpB%d", i), []byte("B")),
   106  				)
   107  			}
   108  		}
   109  		router.Register(kit.NewRouter(&OpA{}))
   110  		router.Register(kit.NewRouter(&OpB{}))
   111  
   112  		tw := mq.NewTaskWorker(tm,
   113  			mq.WithChannel(channel),
   114  			mq.WithWorkerCount(2),
   115  			mq.WithFinishFunc(func(ctx context.Context, task mq.Task) {
   116  				NewWithT(t).Expect(task.State()).To(Equal(mq.TASK_STATE__SUCCEEDED))
   117  			}),
   118  		)
   119  
   120  		go func() {
   121  			fmt.Println(tw.Serve(router))
   122  		}()
   123  
   124  		time.Sleep(800 * time.Microsecond)
   125  
   126  		p, _ := os.FindProcess(os.Getpid())
   127  		_ = p.Signal(os.Interrupt)
   128  
   129  		time.Sleep(time.Second)
   130  	}
   131  }