forked from iovisor/bcc
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_stat1.py
executable file
·60 lines (53 loc) · 1.83 KB
/
test_stat1.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#!/usr/bin/env python
# Copyright (c) PLUMgrid, Inc.
# Licensed under the Apache License, Version 2.0 (the "License")
# test program to count the packets sent to a device in a .5
# second period
from ctypes import c_uint, c_ulong, Structure
from netaddr import IPAddress
from bcc import BPF
from subprocess import check_call
import sys
from unittest import main, TestCase
arg1 = sys.argv.pop(1)
arg2 = ""
if len(sys.argv) > 1:
arg2 = sys.argv.pop(1)
Key = None
Leaf = None
if arg1.endswith(".b"):
class Key(Structure):
_fields_ = [("dip", c_uint),
("sip", c_uint)]
class Leaf(Structure):
_fields_ = [("rx_pkts", c_ulong),
("tx_pkts", c_ulong)]
class TestBPFSocket(TestCase):
def setUp(self):
b = BPF(arg1, arg2, debug=0)
fn = b.load_func("on_packet", BPF.SOCKET_FILTER)
BPF.attach_raw_socket(fn, "eth0")
self.stats = b.get_table("stats", Key, Leaf)
def test_ping(self):
cmd = ["ping", "-f", "-c", "100", "172.16.1.1"]
check_call(cmd)
#for key, leaf in self.stats.items():
# print(IPAddress(key.sip), "=>", IPAddress(key.dip),
# "rx", leaf.rx_pkts, "tx", leaf.tx_pkts)
key = self.stats.Key(IPAddress("172.16.1.2").value, IPAddress("172.16.1.1").value)
leaf = self.stats[key]
self.assertEqual(leaf.rx_pkts, 100)
self.assertEqual(leaf.tx_pkts, 100)
del self.stats[key]
with self.assertRaises(KeyError):
x = self.stats[key]
with self.assertRaises(KeyError):
del self.stats[key]
self.stats.clear()
self.assertEqual(len(self.stats), 0)
self.stats[key] = leaf
self.assertEqual(len(self.stats), 1)
self.stats.clear()
self.assertEqual(len(self.stats), 0)
if __name__ == "__main__":
main()