From 533391eaa94c4b3a148bf801b8aea5bb803c1e9d Mon Sep 17 00:00:00 2001
From: Yonghong Song <yhs@fb.com>
Date: Sun, 5 Apr 2020 22:21:21 -0700
Subject: [PATCH] add test_lpm_trie.py test

test_lpm_trie.py has been in the repo for quite some time,
but is not included in the unit test.

The issue https://github.com/iovisor/bcc/issues/2860
exposed an issue involved in using together with BTF,
which requires the key type to be a structure.

Let add it as a unit test so we can be sure
lpm_trie map is supported properly by bcc.

Signed-off-by: Yonghong Song <yhs@fb.com>
---
 tests/python/CMakeLists.txt   |  2 ++
 tests/python/test_lpm_trie.py | 26 ++++++++++++++++++++++----
 2 files changed, 24 insertions(+), 4 deletions(-)
 mode change 100644 => 100755 tests/python/test_lpm_trie.py

diff --git a/tests/python/CMakeLists.txt b/tests/python/CMakeLists.txt
index f323ac1b334b..63a34130572f 100644
--- a/tests/python/CMakeLists.txt
+++ b/tests/python/CMakeLists.txt
@@ -85,3 +85,5 @@ add_test(NAME py_test_free_bcc_memory WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_D
   COMMAND ${TEST_WRAPPER} py_test_free_bcc_memory sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_free_bcc_memory.py)
 add_test(NAME py_test_rlimit WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_test_rlimit sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_rlimit.py)
+add_test(NAME py_test_lpm_trie WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_test_lpm_trie sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_lpm_trie.py)
diff --git a/tests/python/test_lpm_trie.py b/tests/python/test_lpm_trie.py
old mode 100644
new mode 100755
index 560cb4b6cae4..c95b9ceddd24
--- a/tests/python/test_lpm_trie.py
+++ b/tests/python/test_lpm_trie.py
@@ -3,10 +3,23 @@
 # Licensed under the Apache License, Version 2.0 (the "License")
 
 import ctypes as ct
-import unittest
+import distutils.version
+import os
+from unittest import main, skipUnless, TestCase
 from bcc import BPF
 from netaddr import IPAddress
 
+def kernel_version_ge(major, minor):
+    # True if running kernel is >= X.Y
+    version = distutils.version.LooseVersion(os.uname()[2]).version
+    if version[0] > major:
+        return True
+    if version[0] < major:
+        return False
+    if minor and version[1] < minor:
+        return False
+    return True
+
 class KeyV4(ct.Structure):
     _fields_ = [("prefixlen", ct.c_uint),
                 ("data", ct.c_ubyte * 4)]
@@ -15,10 +28,15 @@ class KeyV6(ct.Structure):
     _fields_ = [("prefixlen", ct.c_uint),
                 ("data", ct.c_ushort * 8)]
 
-class TestLpmTrie(unittest.TestCase):
+@skipUnless(kernel_version_ge(4, 11), "requires kernel >= 4.11")
+class TestLpmTrie(TestCase):
     def test_lpm_trie_v4(self):
         test_prog1 = """
-        BPF_LPM_TRIE(trie, u64, int, 16);
+        struct key_v4 {
+            u32 prefixlen;
+            u32 data[4];
+        };
+        BPF_LPM_TRIE(trie, struct key_v4, int, 16);
         """
         b = BPF(text=test_prog1)
         t = b["trie"]
@@ -71,4 +89,4 @@ def test_lpm_trie_v6(self):
             v = t[k]
 
 if __name__ == "__main__":
-    unittest.main()
+    main()