kraken

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
}
96