github.com/unigraph-dev/dgraph@v1.1.1-0.20200923154953-8b52b426f765/conn/node_test.go (about)

     1  /*
     2   * Copyright 2018 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package conn
    18  
    19  import (
    20  	"bytes"
    21  	"fmt"
    22  	"io/ioutil"
    23  	"os"
    24  	"sync"
    25  	"testing"
    26  	"time"
    27  
    28  	"github.com/dgraph-io/badger"
    29  	"github.com/dgraph-io/dgraph/protos/pb"
    30  	"github.com/dgraph-io/dgraph/raftwal"
    31  	"github.com/stretchr/testify/require"
    32  	"go.etcd.io/etcd/raft"
    33  	"go.etcd.io/etcd/raft/raftpb"
    34  	"golang.org/x/net/context"
    35  )
    36  
    37  func (n *Node) run(wg *sync.WaitGroup) {
    38  	ticker := time.NewTicker(20 * time.Millisecond)
    39  	defer ticker.Stop()
    40  
    41  	for {
    42  		select {
    43  		case <-ticker.C:
    44  			n.Raft().Tick()
    45  		case rd := <-n.Raft().Ready():
    46  			n.SaveToStorage(rd.HardState, rd.Entries, rd.Snapshot)
    47  			for _, entry := range rd.CommittedEntries {
    48  				if entry.Type == raftpb.EntryConfChange {
    49  					var cc raftpb.ConfChange
    50  					cc.Unmarshal(entry.Data)
    51  					n.Raft().ApplyConfChange(cc)
    52  				} else if entry.Type == raftpb.EntryNormal {
    53  					if bytes.HasPrefix(entry.Data, []byte("hey")) {
    54  						wg.Done()
    55  					}
    56  				}
    57  			}
    58  			n.Raft().Advance()
    59  		}
    60  	}
    61  }
    62  
    63  func TestProposal(t *testing.T) {
    64  	dir, err := ioutil.TempDir("", "badger")
    65  	require.NoError(t, err)
    66  	defer os.RemoveAll(dir)
    67  
    68  	db, err := badger.Open(badger.DefaultOptions(dir))
    69  	require.NoError(t, err)
    70  	store := raftwal.Init(db, 0, 0)
    71  
    72  	rc := &pb.RaftContext{Id: 1}
    73  	n := NewNode(rc, store)
    74  
    75  	peers := []raft.Peer{{ID: n.Id}}
    76  	n.SetRaft(raft.StartNode(n.Cfg, peers))
    77  
    78  	loop := 5
    79  	var wg sync.WaitGroup
    80  	wg.Add(loop)
    81  	go n.run(&wg)
    82  
    83  	for i := 0; i < loop; i++ {
    84  		data := []byte(fmt.Sprintf("hey-%d", i))
    85  		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
    86  		defer cancel()
    87  		require.NoError(t, n.Raft().Propose(ctx, data))
    88  	}
    89  	wg.Wait()
    90  }