youtubebeat/vendor/github.com/elastic/beats/packetbeat/flows/flows_test.go

161 lines
3.9 KiB
Go
Raw Normal View History

2018-11-18 11:08:38 +01:00
// Licensed to Elasticsearch B.V. under one or more contributor
// license agreements. See the NOTICE file distributed with
// this work for additional information regarding copyright
// ownership. Elasticsearch B.V. licenses this file to you 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.
// +build !integration
package flows
import (
"net"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/elastic/beats/libbeat/beat"
"github.com/elastic/beats/libbeat/common"
"github.com/elastic/beats/libbeat/logp"
"github.com/elastic/beats/packetbeat/config"
)
type flowsChan struct {
ch chan []beat.Event
}
func (f *flowsChan) PublishFlows(events []beat.Event) {
f.ch <- events
}
func TestFlowsCounting(t *testing.T) {
logp.TestingSetup()
mac1 := []byte{1, 2, 3, 4, 5, 6}
mac2 := []byte{6, 5, 4, 3, 2, 1}
ip1 := []byte{127, 0, 0, 1}
ip2 := []byte{128, 0, 1, 2}
port1 := []byte{0, 1}
port2 := []byte{0, 2}
module, err := NewFlows(nil, &config.Flows{})
assert.NoError(t, err)
uint1, err := module.NewUint("uint1")
uint2, err := module.NewUint("uint2")
int1, err := module.NewInt("int1")
int2, err := module.NewInt("int2")
float1, err := module.NewFloat("float1")
float2, err := module.NewFloat("float2")
assert.NoError(t, err)
pub := &flowsChan{make(chan []beat.Event, 1)}
processor := &flowsProcessor{
table: module.table,
counters: module.counterReg,
timeout: 20 * time.Millisecond,
}
processor.spool.init(pub.PublishFlows, 1)
worker, err := makeWorker(
processor,
10*time.Millisecond,
1,
-1,
0)
if err != nil {
t.Fatalf("Failed to create flow worker: %v", err)
}
worker.Start()
defer worker.Stop()
idForward := newFlowID()
addrForward := addAll(
addEther(mac1, mac2),
addIP(ip1, ip2),
addTCP(port1, port2),
)
addrForward(idForward)
idRev := newFlowID()
addrRev := addAll(
addEther(mac2, mac1),
addIP(ip2, ip1),
addTCP(port2, port1),
)
addrRev(idRev)
assert.True(t, FlowIDsEqual(idForward, idRev))
{
module.Lock()
flow := module.Get(idForward)
flowRev := module.Get(idRev)
int1.Add(flow, -1)
uint1.Add(flow, 1)
float1.Add(flow, 3.14)
int2.Set(flowRev, -1)
uint2.Set(flowRev, 5)
float2.Set(flowRev, 1.4142)
module.Unlock()
}
var events []beat.Event
select {
case events = <-pub.ch:
case <-time.After(5 * time.Second):
}
if events == nil {
t.Fatalf("no event received in time")
}
event := events[0].Fields
t.Logf("event: %v", event)
source := event["source"].(common.MapStr)
dest := event["dest"].(common.MapStr)
// validate generated event
assert.Equal(t, net.HardwareAddr(mac1).String(), source["mac"])
assert.Equal(t, net.HardwareAddr(mac2).String(), dest["mac"])
assert.Equal(t, net.IP(ip1).String(), source["ip"])
assert.Equal(t, net.IP(ip2).String(), dest["ip"])
assert.Equal(t, uint16(256), source["port"])
assert.Equal(t, uint16(512), dest["port"])
assert.Equal(t, "tcp", event["transport"])
stat := source["stats"].(map[string]interface{})
assert.Equal(t, int64(-1), stat["int1"])
assert.Equal(t, nil, stat["int2"])
assert.Equal(t, uint64(1), stat["uint1"])
assert.Equal(t, nil, stat["uint2"])
assert.Equal(t, 3.14, stat["float1"])
assert.Equal(t, nil, stat["float2"])
stat = dest["stats"].(map[string]interface{})
assert.Equal(t, nil, stat["int1"])
assert.Equal(t, int64(-1), stat["int2"])
assert.Equal(t, nil, stat["uint1"])
assert.Equal(t, uint64(5), stat["uint2"])
assert.Equal(t, nil, stat["float1"])
assert.Equal(t, 1.4142, stat["float2"])
}