orbit/pulse/service/batch_test.go

package service

import (
	"context"
	"encoding/json"
	"errors"
	"strings"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/orbit/internal/fault"
	core "git.yuntuo.hk/k2safe/orbit/model"
	"git.yuntuo.hk/k2safe/orbit/platform/auth"
)

type batchStore struct {
	Store
	authErr          error
	authCalls, calls int
	submit           func(int, string, SubmitInput) (SubmitResult, error)
}

func (s *batchStore) AuthorizeSubmit(context.Context, auth.Token) error {
	s.authCalls++
	return s.authErr
}
func (s *batchStore) Submit(_ context.Context, _ auth.Token, key string, in SubmitInput) (SubmitResult, error) {
	s.calls++
	return s.submit(s.calls, key, in)
}
func batchInput(n int) BatchSubmitInput {
	in := BatchSubmitInput{}
	for i := 0; i < n; i++ {
		in.Items = append(in.Items, BatchSubmitItem{IdempotencyKey: string(rune('a' + i)), Task: SubmitInput{QueueID: core.ID(strings.Repeat("1", 32)), Handler: "file.convert", HandlerVersion: "v1", PayloadVersion: "v1", Payload: json.RawMessage(`{"n":9007199254740993}`)}})
	}
	return in
}
func acceptedBatchTask(in SubmitInput) SubmitResult {
	return SubmitResult{Task: TaskView{ID: core.ID(strings.Repeat("2", 32)), QueueID: in.QueueID, Revision: 1, CreatedAt: time.Now().UTC()}}
}

func TestBatchEnvelopeBeforeWritesAndAuthorization(t *testing.T) {
	for _, n := range []int{0, 33} {
		s := &batchStore{}
		if _, err := New(s).SubmitBatch(context.Background(), auth.Token{}, batchInput(n)); fault.Code(err) != "INVALID_ARGUMENT" || s.authCalls != 0 || s.calls != 0 {
			t.Fatal(n, err)
		}
	}
	for _, key := range []string{"", "a b", strings.Repeat("a", 129), "a"} {
		in := batchInput(2)
		in.Items[1].IdempotencyKey = key
		s := &batchStore{}
		if _, err := New(s).SubmitBatch(context.Background(), auth.Token{}, in); fault.Code(err) != "INVALID_ARGUMENT" || s.calls != 0 {
			t.Fatal(key, err)
		}
	}
	for _, denied := range []error{fault.Unauthorized, fault.Forbidden, errors.New("database private detail")} {
		s := &batchStore{authErr: denied}
		in := batchInput(1)
		in.Items[0].Task.Handler = ""
		if _, err := New(s).SubmitBatch(context.Background(), auth.Token{}, in); err != denied || s.calls != 0 || s.authCalls != 1 {
			t.Fatal(err)
		}
	}
}

func TestBatchPartialResultsAndIndependentNormalization(t *testing.T) {
	in := batchInput(5)
	in.Items[1].Task.Priority = 10
	s := &batchStore{submit: func(n int, key string, in SubmitInput) (SubmitResult, error) {
		if in.TaskTimeoutMS == 0 || in.Retry.MaxAttempts != 1 || string(in.Payload) != `{"n":9007199254740993}` {
			t.Fatal("normalization or precision lost")
		}
		switch key {
		case "c":
			return SubmitResult{}, fault.QueueFull
		case "d":
			return SubmitResult{}, fault.IdempotencyConflict
		}
		out := acceptedBatchTask(in)
		out.Deduplicated = key == "e"
		return out, nil
	}}
	out, err := New(s).SubmitBatch(context.Background(), auth.Token{}, in)
	if err != nil || len(out.Items) != 5 || s.calls != 4 {
		t.Fatal(out, err)
	}
	for i, want := range []BatchItemStatus{BatchAccepted, BatchRejected, BatchRejected, BatchRejected, BatchAccepted} {
		if out.Items[i].Status != want || out.Items[i].Index != i {
			t.Fatal(out)
		}
	}
	if out.Items[1].Error.Code != "INVALID_ARGUMENT" || out.Items[2].Error.RetryAfterMS != 1000 || !out.Items[4].Receipt.Deduplicated {
		t.Fatal(out)
	}
}

func TestBatchUnknownAndCancellationNeverClaimRollback(t *testing.T) {
	for _, fail := range []error{errors.New("secret database error"), context.DeadlineExceeded, fault.Unavailable, fault.Unauthorized} {
		s := &batchStore{submit: func(n int, _ string, in SubmitInput) (SubmitResult, error) {
			if n == 1 {
				return acceptedBatchTask(in), nil
			}
			return SubmitResult{}, fail
		}}
		out, err := New(s).SubmitBatch(context.Background(), auth.Token{}, batchInput(4))
		if err != nil || s.calls != 2 || out.Items[0].Status != BatchAccepted || out.Items[2].Status != BatchNotAttempted || out.Items[3].Status != BatchNotAttempted {
			t.Fatal(out, err)
		}
		want := BatchUnknown
		if fail == fault.Unauthorized {
			want = BatchRejected
		}
		if out.Items[1].Status != want || strings.Contains(out.Items[1].Error.Message, "secret") {
			t.Fatal(out)
		}
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	s := &batchStore{submit: func(_ int, _ string, in SubmitInput) (SubmitResult, error) {
		cancel()
		return acceptedBatchTask(in), nil
	}}
	out, err := New(s).SubmitBatch(ctx, auth.Token{}, batchInput(3))
	if err != nil || s.calls != 1 || out.Items[0].Status != BatchAccepted || out.Items[1].Status != BatchNotAttempted {
		t.Fatal(out, err)
	}
}