github.com/songzhibin97/gkit@v1.2.13/distributed/backend/backend_mongodb/mongo_test.go (about)

     1  package backend_mongodb
     2  
     3  import (
     4  	"context"
     5  	"strconv"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/songzhibin97/gkit/generator"
    10  
    11  	"github.com/stretchr/testify/assert"
    12  
    13  	"github.com/songzhibin97/gkit/distributed/backend"
    14  	"github.com/songzhibin97/gkit/distributed/task"
    15  	"go.mongodb.org/mongo-driver/mongo"
    16  )
    17  
    18  func InitBackend() backend.Backend {
    19  	client, err := mongo.NewClient()
    20  	if err != nil {
    21  		return nil
    22  	}
    23  	err = client.Connect(context.Background())
    24  	if err != nil {
    25  		return nil
    26  	}
    27  	return NewBackendMongoDB(client, -1)
    28  }
    29  
    30  func TestGroupTaskOver(t *testing.T) {
    31  	_backend := InitBackend()
    32  	if _backend == nil {
    33  		t.Skip()
    34  	}
    35  	g := generator.NewSnowflake(time.Now().Local(), 1)
    36  	ids := make([]string, 0, 3)
    37  	for i := 0; i < 3; i++ {
    38  		id, _ := g.NextID()
    39  		if i == 0 {
    40  			ids = append(ids, "group:"+strconv.FormatUint(id, 10))
    41  		} else {
    42  			ids = append(ids, "task:"+strconv.FormatUint(id, 10))
    43  		}
    44  
    45  	}
    46  	var (
    47  		group = task.GroupMeta{
    48  			GroupID: ids[0],
    49  			Name:    "group",
    50  		}
    51  		task1 = task.Signature{
    52  			ID:      ids[1],
    53  			GroupID: group.GroupID,
    54  			Name:    "task1",
    55  		}
    56  		task2 = task.Signature{
    57  			ID:      ids[2],
    58  			GroupID: group.GroupID,
    59  			Name:    "task2",
    60  		}
    61  	)
    62  	_ = _backend.ResetGroup(group.GroupID)
    63  	_ = _backend.ResetTask(task1.ID, task2.ID)
    64  	isCompleted, err := _backend.GroupCompleted(group.GroupID)
    65  	if assert.Error(t, err) {
    66  		assert.False(t, isCompleted)
    67  		assert.Error(t, err, mongo.ErrNoDocuments)
    68  	}
    69  	_ = _backend.GroupTakeOver(group.GroupID, group.Name, task1.ID, task2.ID)
    70  	isCompleted, err = _backend.GroupCompleted(group.GroupID)
    71  	if assert.NoError(t, err) {
    72  		assert.False(t, isCompleted)
    73  	}
    74  
    75  	_ = _backend.SetStatePending(&task1)
    76  	_ = _backend.SetStateStarted(&task2)
    77  	isCompleted, err = _backend.GroupCompleted(group.GroupID)
    78  	if assert.NoError(t, err) {
    79  		assert.False(t, isCompleted)
    80  	}
    81  	result := []*task.Result{{
    82  		Type:  "int",
    83  		Value: 1,
    84  	}}
    85  	_ = _backend.SetStateStarted(&task1)
    86  	_ = _backend.SetStateSuccess(&task2, result)
    87  	isCompleted, err = _backend.GroupCompleted(group.GroupID)
    88  	if assert.NoError(t, err) {
    89  		assert.False(t, isCompleted)
    90  	}
    91  	_ = _backend.SetStateFailure(&task1, "failure")
    92  	isCompleted, err = _backend.GroupCompleted(group.GroupID)
    93  	if assert.NoError(t, err) {
    94  		assert.True(t, isCompleted)
    95  	}
    96  }
    97  
    98  func TestGetStatus(t *testing.T) {
    99  	_backend := InitBackend()
   100  	if _backend == nil {
   101  		t.Skip()
   102  	}
   103  	task1 := task.Signature{
   104  		ID:      "task1",
   105  		GroupID: "group",
   106  		Name:    "task1",
   107  	}
   108  	_ = _backend.ResetTask(task1.ID)
   109  
   110  	status, err := _backend.GetStatus(task1.ID)
   111  	assert.Equal(t, err, mongo.ErrNoDocuments)
   112  	assert.Nil(t, status)
   113  
   114  	_ = _backend.SetStatePending(&task1)
   115  	status, err = _backend.GetStatus(task1.ID)
   116  	assert.NoError(t, err)
   117  	assert.Equal(t, status.Status, task.StatePending)
   118  
   119  	_ = _backend.SetStateReceived(&task1)
   120  	status, err = _backend.GetStatus(task1.ID)
   121  	assert.NoError(t, err)
   122  	assert.Equal(t, status.Status, task.StateReceived)
   123  
   124  	_ = _backend.SetStateStarted(&task1)
   125  	status, err = _backend.GetStatus(task1.ID)
   126  	assert.NoError(t, err)
   127  	assert.Equal(t, status.Status, task.StateStarted)
   128  
   129  	result := &task.Result{
   130  		Type:  "int",
   131  		Value: 1,
   132  	}
   133  	_ = _backend.SetStateSuccess(&task1, []*task.Result{result})
   134  	status, err = _backend.GetStatus(task1.ID)
   135  	assert.NoError(t, err)
   136  	assert.Equal(t, status.Status, task.StateSuccess)
   137  }
   138  
   139  func TestResult(t *testing.T) {
   140  	_backend := InitBackend()
   141  	if _backend == nil {
   142  		t.Skip()
   143  	}
   144  	task1 := task.Signature{
   145  		ID:      "task1",
   146  		GroupID: "group",
   147  		Name:    "task1",
   148  	}
   149  	_ = _backend.SetStatePending(&task1)
   150  	status, err := _backend.GetStatus(task1.ID)
   151  	assert.NoError(t, err)
   152  	assert.Equal(t, status.Status, task.StatePending)
   153  
   154  	_ = _backend.ResetTask(task1.ID)
   155  
   156  	status, err = _backend.GetStatus(task1.ID)
   157  	assert.Equal(t, err, mongo.ErrNoDocuments)
   158  	assert.Nil(t, status)
   159  }