go.etcd.io/etcd@v3.3.27+incompatible/rafthttp/snapshot_test.go (about)

     1  // Copyright 2016 The etcd Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package rafthttp
    16  
    17  import (
    18  	"fmt"
    19  	"io"
    20  	"io/ioutil"
    21  	"net/http"
    22  	"net/http/httptest"
    23  	"os"
    24  	"strings"
    25  	"testing"
    26  	"time"
    27  
    28  	"github.com/coreos/etcd/pkg/types"
    29  	"github.com/coreos/etcd/raft/raftpb"
    30  	"github.com/coreos/etcd/snap"
    31  )
    32  
    33  type strReaderCloser struct{ *strings.Reader }
    34  
    35  func (s strReaderCloser) Close() error { return nil }
    36  
    37  func TestSnapshotSend(t *testing.T) {
    38  	tests := []struct {
    39  		m    raftpb.Message
    40  		rc   io.ReadCloser
    41  		size int64
    42  
    43  		wsent  bool
    44  		wfiles int
    45  	}{
    46  		// sent and receive with no errors
    47  		{
    48  			m:    raftpb.Message{Type: raftpb.MsgSnap, To: 1},
    49  			rc:   strReaderCloser{strings.NewReader("hello")},
    50  			size: 5,
    51  
    52  			wsent:  true,
    53  			wfiles: 1,
    54  		},
    55  		// error when reading snapshot for send
    56  		{
    57  			m:    raftpb.Message{Type: raftpb.MsgSnap, To: 1},
    58  			rc:   &errReadCloser{fmt.Errorf("snapshot error")},
    59  			size: 1,
    60  
    61  			wsent:  false,
    62  			wfiles: 0,
    63  		},
    64  		// sends less than the given snapshot length
    65  		{
    66  			m:    raftpb.Message{Type: raftpb.MsgSnap, To: 1},
    67  			rc:   strReaderCloser{strings.NewReader("hello")},
    68  			size: 10000,
    69  
    70  			wsent:  false,
    71  			wfiles: 0,
    72  		},
    73  		// sends less than actual snapshot length
    74  		{
    75  			m:    raftpb.Message{Type: raftpb.MsgSnap, To: 1},
    76  			rc:   strReaderCloser{strings.NewReader("hello")},
    77  			size: 1,
    78  
    79  			wsent:  false,
    80  			wfiles: 0,
    81  		},
    82  	}
    83  
    84  	for i, tt := range tests {
    85  		sent, files := testSnapshotSend(t, snap.NewMessage(tt.m, tt.rc, tt.size))
    86  		if tt.wsent != sent {
    87  			t.Errorf("#%d: snapshot expected %v, got %v", i, tt.wsent, sent)
    88  		}
    89  		if tt.wfiles != len(files) {
    90  			t.Fatalf("#%d: expected %d files, got %d files", i, tt.wfiles, len(files))
    91  		}
    92  	}
    93  }
    94  
    95  func testSnapshotSend(t *testing.T, sm *snap.Message) (bool, []os.FileInfo) {
    96  	d, err := ioutil.TempDir(os.TempDir(), "snapdir")
    97  	if err != nil {
    98  		t.Fatal(err)
    99  	}
   100  	defer os.RemoveAll(d)
   101  
   102  	r := &fakeRaft{}
   103  	tr := &Transport{pipelineRt: &http.Transport{}, ClusterID: types.ID(1), Raft: r}
   104  	ch := make(chan struct{}, 1)
   105  	h := &syncHandler{newSnapshotHandler(tr, r, snap.New(d), types.ID(1)), ch}
   106  	srv := httptest.NewServer(h)
   107  	defer srv.Close()
   108  
   109  	picker := mustNewURLPicker(t, []string{srv.URL})
   110  	snapsend := newSnapshotSender(tr, picker, types.ID(1), newPeerStatus(types.ID(1)))
   111  	defer snapsend.stop()
   112  
   113  	snapsend.send(*sm)
   114  
   115  	sent := false
   116  	select {
   117  	case <-time.After(time.Second):
   118  		t.Fatalf("timed out sending snapshot")
   119  	case sent = <-sm.CloseNotify():
   120  	}
   121  
   122  	// wait for handler to finish accepting snapshot
   123  	<-ch
   124  
   125  	files, rerr := ioutil.ReadDir(d)
   126  	if rerr != nil {
   127  		t.Fatal(rerr)
   128  	}
   129  	return sent, files
   130  }
   131  
   132  type errReadCloser struct{ err error }
   133  
   134  func (s *errReadCloser) Read(p []byte) (int, error) { return 0, s.err }
   135  func (s *errReadCloser) Close() error               { return s.err }
   136  
   137  type syncHandler struct {
   138  	h  http.Handler
   139  	ch chan<- struct{}
   140  }
   141  
   142  func (sh *syncHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
   143  	sh.h.ServeHTTP(w, r)
   144  	sh.ch <- struct{}{}
   145  }