tidb/pkg/timer/runtime/main_test.go

// Copyright 2023 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package runtime

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/pingcap/tidb/pkg/testkit/testsetup"
	"github.com/pingcap/tidb/pkg/timer/api"
	"github.com/pingcap/tidb/pkg/util"
	"github.com/stretchr/testify/mock"
	"go.uber.org/goleak"
)

func TestMain(m *testing.M) {
	testsetup.SetupForCommonTest()
	opts := []goleak.Option{
		goleak.IgnoreTopFunction("github.com/golang/glog.(*fileSink).flushDaemon"),
		goleak.IgnoreTopFunction("github.com/bazelbuild/rules_go/go/tools/bzltestutil.RegisterTimeoutHandler.func1"),
		goleak.IgnoreTopFunction("github.com/lestrrat-go/httprc.runFetchWorker"),
		goleak.IgnoreTopFunction("go.opencensus.io/stats/view.(*worker).start"),
	}
	testsetup.SetupForCommonTest()
	goleak.VerifyTestMain(m, opts...)
}

type mockHook struct {
	mock.Mock
	started chan struct{}
	stopped chan struct{}
}

type newHookFn struct {
	mock.Mock
}

func (n *newHookFn) OnFuncCall() *mock.Call {
	return n.On("Func")
}

func (n *newHookFn) Func() api.Hook {
	args := n.Called()
	if v := args.Get(0); v != nil {
		return v.(api.Hook)
	}
	return nil
}

func onlyOnceNewHook(hook api.Hook) func() api.Hook {
	n := newHookFn{}
	n.OnFuncCall().Return(hook).Once()
	return n.Func
}

func newMockHook() *mockHook {
	return &mockHook{
		started: make(chan struct{}),
		stopped: make(chan struct{}),
	}
}

func (h *mockHook) Start() {
	h.Called()
	close(h.started)
}

func (h *mockHook) Stop() {
	h.Called()
	close(h.stopped)
}

func (h *mockHook) OnPreSchedEvent(ctx context.Context, event api.TimerShedEvent) (api.PreSchedEventResult, error) {
	args := h.Called(ctx, event)
	return args.Get(0).(api.PreSchedEventResult), args.Error(1)
}

func (h *mockHook) OnSchedEvent(ctx context.Context, event api.TimerShedEvent) error {
	args := h.Called(ctx, event)
	return args.Error(0)
}

type mockStoreCore struct {
	mock.Mock
}

func newMockStore() (*mockStoreCore, *api.TimerStore) {
	core := &mockStoreCore{}
	return core, &api.TimerStore{TimerStoreCore: core}
}

func (s *mockStoreCore) mock() *mock.Mock {
	return &s.Mock
}

func (s *mockStoreCore) Create(ctx context.Context, record *api.TimerRecord) (string, error) {
	args := s.Called(ctx, record)
	return args.String(0), args.Error(1)
}

func (s *mockStoreCore) List(ctx context.Context, cond api.Cond) ([]*api.TimerRecord, error) {
	args := s.Called(ctx, cond)
	return args.Get(0).([]*api.TimerRecord), args.Error(1)
}

func (s *mockStoreCore) Update(ctx context.Context, timerID string, update *api.TimerUpdate) error {
	args := s.Called(ctx, timerID, update)
	return args.Error(0)
}

func (s *mockStoreCore) Delete(ctx context.Context, timerID string) (bool, error) {
	args := s.Called(ctx, timerID)
	return args.Bool(0), args.Error(1)
}

func (s *mockStoreCore) WatchSupported() bool {
	args := s.Called()
	return args.Bool(0)
}

func (s *mockStoreCore) Watch(ctx context.Context) api.WatchTimerChan {
	args := s.Called(ctx)
	return args.Get(0).(api.WatchTimerChan)
}

func (s *mockStoreCore) Close() {
}

func waitDone(obj any, timeout time.Duration) {
	var ch <-chan struct{}
	switch o := obj.(type) {
	case chan struct{}:
		ch = o
	case <-chan struct{}:
		ch = o
	case *util.WaitGroupWrapper:
		newCh := make(chan struct{})
		ch = newCh

		go func() {
			o.Wait()
			close(newCh)
		}()
	case *sync.WaitGroup:
		newCh := make(chan struct{})
		ch = newCh

		go func() {
			o.Wait()
			close(newCh)
		}()
	default:
		panic(fmt.Sprintf("unsupported type: %T", obj))
	}

	tm := time.NewTimer(timeout)
	defer tm.Stop()
	select {
	case <-ch:
		return
	case <-tm.C:
		panic("wait done timeout")
	}
}

func checkNotDone(ch <-chan struct{}, after time.Duration) {
	if after != 0 {
		time.Sleep(after)
	}
	select {
	case <-ch:
		panic("the channel is expected not done")
	default:
	}
}