// Copyright 2015 The etcd Authors
//
// 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 raft

import (
	"testing"

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

	pb "go.etcd.io/raft/v3/raftpb"
)

var (
	testingSnap = &pb.Snapshot{
		Metadata: &pb.SnapshotMetadata{
			Index:     new(uint64(11)), // magic number
			Term:      new(uint64(11)), // magic number
			ConfState: &pb.ConfState{Voters: []uint64{1, 2}},
		},
	}
)

func TestSendingSnapshotSetPendingSnapshot(t *testing.T) {
	storage := newTestMemoryStorage(withPeers(1))
	sm := newTestRaft(1, 10, 1, storage)
	sm.restore(testingSnap)

	sm.becomeCandidate()
	sm.becomeLeader()

	// force set the next of node 2, so that
	// node 2 needs a snapshot
	sm.trk.Progress[2].Next = sm.raftLog.firstIndex()

	sm.Step(&pb.Message{From: new(uint64(2)), To: new(uint64(1)), Type: pb.MsgAppResp.Enum(), Index: new(sm.trk.Progress[2].Next - 1), Reject: new(true)})
	require.Equal(t, uint64(11), sm.trk.Progress[2].PendingSnapshot)
}

func TestPendingSnapshotPauseReplication(t *testing.T) {
	storage := newTestMemoryStorage(withPeers(1, 2))
	sm := newTestRaft(1, 10, 1, storage)
	sm.restore(testingSnap)

	sm.becomeCandidate()
	sm.becomeLeader()

	sm.trk.Progress[2].BecomeSnapshot(11)

	sm.Step(&pb.Message{From: new(uint64(1)), To: new(uint64(1)), Type: pb.MsgProp.Enum(), Entries: []*pb.Entry{{Data: []byte("somedata")}}})
	msgs := sm.readMessages()
	require.Empty(t, msgs)
}

func TestSnapshotFailure(t *testing.T) {
	storage := newTestMemoryStorage(withPeers(1, 2))
	sm := newTestRaft(1, 10, 1, storage)
	sm.restore(testingSnap)

	sm.becomeCandidate()
	sm.becomeLeader()

	sm.trk.Progress[2].Next = 1
	sm.trk.Progress[2].BecomeSnapshot(11)

	sm.Step(&pb.Message{From: new(uint64(2)), To: new(uint64(1)), Type: pb.MsgSnapStatus.Enum(), Reject: new(true)})
	require.Zero(t, sm.trk.Progress[2].PendingSnapshot)
	require.Equal(t, uint64(1), sm.trk.Progress[2].Next)
	assert.True(t, sm.trk.Progress[2].MsgAppFlowPaused)
}

func TestSnapshotSucceed(t *testing.T) {
	storage := newTestMemoryStorage(withPeers(1, 2))
	sm := newTestRaft(1, 10, 1, storage)
	sm.restore(testingSnap)

	sm.becomeCandidate()
	sm.becomeLeader()

	sm.trk.Progress[2].Next = 1
	sm.trk.Progress[2].BecomeSnapshot(11)

	sm.Step(&pb.Message{From: new(uint64(2)), To: new(uint64(1)), Type: pb.MsgSnapStatus.Enum(), Reject: new(false)})
	require.Zero(t, sm.trk.Progress[2].PendingSnapshot)
	require.Equal(t, uint64(12), sm.trk.Progress[2].Next)
	assert.True(t, sm.trk.Progress[2].MsgAppFlowPaused)
}

func TestSnapshotAbort(t *testing.T) {
	storage := newTestMemoryStorage(withPeers(1, 2))
	sm := newTestRaft(1, 10, 1, storage)
	sm.restore(testingSnap)

	sm.becomeCandidate()
	sm.becomeLeader()

	sm.trk.Progress[2].Next = 1
	sm.trk.Progress[2].BecomeSnapshot(11)

	// A successful msgAppResp that has a higher/equal index than the
	// pending snapshot should abort the pending snapshot.
	sm.Step(&pb.Message{From: new(uint64(2)), To: new(uint64(1)), Type: pb.MsgAppResp.Enum(), Index: new(uint64(11))})
	require.Zero(t, sm.trk.Progress[2].PendingSnapshot)
	// The follower entered StateReplicate and the leader send an append
	// and optimistically updated the progress (so we see 13 instead of 12).
	// There is something to append because the leader appended an empty entry
	// to the log at index 12 when it assumed leadership.
	require.Equal(t, uint64(13), sm.trk.Progress[2].Next)
	require.Equal(t, 1, sm.trk.Progress[2].Inflights.Count())
}