milvus/internal/rootcoord/step_test.go

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 rootcoord

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

func restoreConfirmGCInterval() {
	confirmGCInterval = time.Minute * 20
}

func Test_confirmGCStep_Execute(t *testing.T) {
	t.Run("wait for reschedule", func(t *testing.T) {
		confirmGCInterval = time.Minute * 1000
		defer restoreConfirmGCInterval()

		s := &confirmGCStep{lastScheduledTime: time.Now()}
		_, err := s.Execute(context.TODO())
		assert.Error(t, err)
	})

	t.Run("GC not finished", func(t *testing.T) {
		broker := newMockBroker()
		broker.GCConfirmFunc = func(ctx context.Context, collectionID, partitionID UniqueID) bool {
			return false
		}

		core := newTestCore(withBroker(broker))

		confirmGCInterval = time.Millisecond
		defer restoreConfirmGCInterval()

		s := newConfirmGCStep(core, 100, 1000)
		time.Sleep(confirmGCInterval)

		_, err := s.Execute(context.TODO())
		assert.Error(t, err)
	})

	t.Run("normal case", func(t *testing.T) {
		broker := newMockBroker()
		broker.GCConfirmFunc = func(ctx context.Context, collectionID, partitionID UniqueID) bool {
			return true
		}

		core := newTestCore(withBroker(broker))

		confirmGCInterval = time.Millisecond
		defer restoreConfirmGCInterval()

		s := newConfirmGCStep(core, 100, 1000)
		time.Sleep(confirmGCInterval)

		_, err := s.Execute(context.TODO())
		assert.NoError(t, err)
	})
}