github.com/uber/kraken@v0.1.4/lib/torrent/networkevent/producer_test.go (about)

     1  // Copyright (c) 2016-2019 Uber Technologies, Inc.
     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  package networkevent
    15  
    16  import (
    17  	"bufio"
    18  	"encoding/json"
    19  	"io/ioutil"
    20  	"os"
    21  	"path/filepath"
    22  	"testing"
    23  
    24  	"github.com/uber/kraken/core"
    25  
    26  	"github.com/stretchr/testify/require"
    27  )
    28  
    29  func TestProducerCreatesAndReusesFile(t *testing.T) {
    30  	require := require.New(t)
    31  
    32  	h := core.InfoHashFixture()
    33  	peer1 := core.PeerIDFixture()
    34  	peer2 := core.PeerIDFixture()
    35  
    36  	dir, err := ioutil.TempDir("", "")
    37  	require.NoError(err)
    38  	defer os.RemoveAll(dir)
    39  
    40  	config := Config{
    41  		Enabled: true,
    42  		LogPath: filepath.Join(dir, "netevents"),
    43  	}
    44  
    45  	events := []*Event{
    46  		ReceivePieceEvent(h, peer1, peer2, 1),
    47  		ReceivePieceEvent(h, peer1, peer2, 2),
    48  		ReceivePieceEvent(h, peer1, peer2, 3),
    49  		ReceivePieceEvent(h, peer1, peer2, 4),
    50  	}
    51  
    52  	// First producer should create the file.
    53  	p, err := NewProducer(config)
    54  	require.NoError(err)
    55  	for _, e := range events[:2] {
    56  		p.Produce(e)
    57  	}
    58  	require.NoError(p.Close())
    59  
    60  	// Second producer should reuse the existing file.
    61  	p, err = NewProducer(config)
    62  	require.NoError(err)
    63  	for _, e := range events[2:] {
    64  		p.Produce(e)
    65  	}
    66  	require.NoError(p.Close())
    67  
    68  	f, err := os.Open(config.LogPath)
    69  	require.NoError(err)
    70  	defer f.Close()
    71  
    72  	var results []*Event
    73  	s := bufio.NewScanner(f)
    74  	s.Split(bufio.ScanLines)
    75  	for s.Scan() {
    76  		e := new(Event)
    77  		require.NoError(json.Unmarshal(s.Bytes(), e))
    78  		results = append(results, e)
    79  	}
    80  
    81  	require.Equal(StripTimestamps(events), StripTimestamps(results))
    82  }
    83  
    84  func TestDisabledProducerNoops(t *testing.T) {
    85  	require := require.New(t)
    86  
    87  	h := core.InfoHashFixture()
    88  	peer1 := core.PeerIDFixture()
    89  	peer2 := core.PeerIDFixture()
    90  
    91  	p, err := NewProducer(Config{})
    92  	require.NoError(err)
    93  
    94  	p.Produce(ReceivePieceEvent(h, peer1, peer2, 1))
    95  }