github.com/iikira/iikira-go-utils@v0.0.0-20230610031953-f2cb11cde33a/utils/taskframework/taskframework_test.go (about)

     1  package taskframework
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  	"time"
     7  )
     8  
     9  type (
    10  	TestUnit struct {
    11  		retry    bool
    12  		taskInfo *TaskInfo
    13  	}
    14  )
    15  
    16  func (tu *TestUnit) SetTaskInfo(taskInfo *TaskInfo) {
    17  	tu.taskInfo = taskInfo
    18  }
    19  
    20  func (tu *TestUnit) OnFailed(lastRunResult *TaskUnitRunResult) {
    21  	fmt.Printf("[%s] error: %s, failed\n", tu.taskInfo.Id(), lastRunResult.Err)
    22  }
    23  
    24  func (tu *TestUnit) OnSuccess(lastRunResult *TaskUnitRunResult) {
    25  	fmt.Printf("[%s] success\n", tu.taskInfo.Id())
    26  }
    27  
    28  func (tu *TestUnit) OnComplete(lastRunResult *TaskUnitRunResult) {
    29  	fmt.Printf("[%s] complete\n", tu.taskInfo.Id())
    30  }
    31  
    32  func (tu *TestUnit) Run() (result *TaskUnitRunResult) {
    33  	fmt.Printf("[%s] running...\n", tu.taskInfo.Id())
    34  	return &TaskUnitRunResult{
    35  		//Succeed:   true,
    36  		NeedRetry: true,
    37  	}
    38  }
    39  
    40  func (tu *TestUnit) OnRetry(lastRunResult *TaskUnitRunResult) {
    41  	fmt.Printf("[%s] prepare retry, times [%d/%d]...\n", tu.taskInfo.Id(), tu.taskInfo.Retry(), tu.taskInfo.MaxRetry())
    42  }
    43  
    44  func (tu *TestUnit) RetryWait() time.Duration {
    45  	return 1 * time.Second
    46  }
    47  
    48  func TestTaskExecutor(t *testing.T) {
    49  	te := NewTaskExecutor()
    50  	te.SetParallel(2)
    51  	for i := 0; i < 3; i++ {
    52  		tu := TestUnit{
    53  			retry: false,
    54  		}
    55  		te.Append(&tu, 2)
    56  	}
    57  	te.Execute()
    58  }