From c6fae35cfdfcde60e1948e876de71674ebaed7f3 Mon Sep 17 00:00:00 2001
From: Junyu Lai <junyulai@google.com>
Date: Thu, 11 Jul 2024 15:57:00 +0800
Subject: [PATCH] Support get apf capabilities from network_stack shell command

This utility reads output of cmd network_stack apf ${iface_name}
capabilities. This is needed for the tests to skip if the apf
version is too old.

Test: atest NetworkStaticLibHostPythonTests
Bug: 352132428
Change-Id: I1f8b523c3e6f687c2e919993bc18392dc25dbbf7
---
 .../tests/unit/host/python/apf_utils_test.py  | 25 +++++++-
 staticlibs/testutils/host/python/apf_utils.py | 64 +++++++++++++++++++
 2 files changed, 88 insertions(+), 1 deletion(-)

diff --git a/staticlibs/tests/unit/host/python/apf_utils_test.py b/staticlibs/tests/unit/host/python/apf_utils_test.py
index 8b390e3fdf..caaf959523 100644
--- a/staticlibs/tests/unit/host/python/apf_utils_test.py
+++ b/staticlibs/tests/unit/host/python/apf_utils_test.py
@@ -13,13 +13,16 @@
 #  limitations under the License.
 
 from unittest.mock import MagicMock, patch
+from absl.testing import parameterized
 from mobly import asserts
 from mobly import base_test
 from mobly import config_parser
 from mobly.controllers.android_device_lib.adb import AdbError
 from net_tests_utils.host.python.apf_utils import (
+    ApfCapabilities,
     PatternNotFoundException,
     UnsupportedOperationException,
+    get_apf_capabilities,
     get_apf_counter,
     get_apf_counters_from_dumpsys,
     get_hardware_address,
@@ -29,7 +32,7 @@ from net_tests_utils.host.python.apf_utils import (
 from net_tests_utils.host.python.assert_utils import UnexpectedBehaviorError
 
 
-class TestApfUtils(base_test.BaseTestClass):
+class TestApfUtils(base_test.BaseTestClass, parameterized.TestCase):
 
   def __init__(self, configs: config_parser.TestRunConfig):
     super().__init__(configs)
@@ -150,3 +153,23 @@ IpClient.wlan1
     )
     with asserts.assert_raises(UnsupportedOperationException):
       send_raw_packet_downstream(self.mock_ad, "eth0", "AABBCCDDEEFF")
+
+  @parameterized.parameters(
+      ("2,2048,1", ApfCapabilities(2, 2048, 1)),  # Valid input
+      ("3,1024,0", ApfCapabilities(3, 1024, 0)),  # Valid input
+      ("invalid,output", ApfCapabilities(0, 0, 0)),  # Invalid input
+      ("", ApfCapabilities(0, 0, 0)),  # Empty input
+  )
+  @patch("net_tests_utils.host.python.adb_utils.adb_shell")
+  def test_get_apf_capabilities(
+      self, mock_output, expected_result, mock_adb_shell
+  ):
+    """Tests the get_apf_capabilities function with various inputs and expected results."""
+    # Configure the mock adb_shell to return the specified output
+    mock_adb_shell.return_value = mock_output
+
+    # Call the function under test
+    result = get_apf_capabilities(self.mock_ad, "wlan0")
+
+    # Assert that the result matches the expected result
+    asserts.assert_equal(result, expected_result)
diff --git a/staticlibs/testutils/host/python/apf_utils.py b/staticlibs/testutils/host/python/apf_utils.py
index f71464cad4..415799c2f6 100644
--- a/staticlibs/testutils/host/python/apf_utils.py
+++ b/staticlibs/testutils/host/python/apf_utils.py
@@ -12,7 +12,9 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from dataclasses import dataclass
 import re
+from mobly import asserts
 from mobly.controllers import android_device
 from mobly.controllers.android_device_lib.adb import AdbError
 from net_tests_utils.host.python import adb_utils, assert_utils
@@ -190,3 +192,65 @@ def send_raw_packet_downstream(
     raise assert_utils.UnexpectedBehaviorError(
         f"Got unexpected output: {output} for command: {cmd}."
     )
+
+
+@dataclass
+class ApfCapabilities:
+  """APF program support capabilities.
+
+  See android.net.apf.ApfCapabilities.
+
+  Attributes:
+      apf_version_supported (int): Version of APF instruction set supported for
+        packet filtering. 0 indicates no support for packet filtering using APF
+        programs.
+      apf_ram_size (int): Size of APF ram.
+      apf_packet_format (int): Format of packets passed to APF filter. Should be
+        one of ARPHRD_*
+  """
+
+  apf_version_supported: int
+  apf_ram_size: int
+  apf_packet_format: int
+
+  def __init__(
+      self,
+      apf_version_supported: int,
+      apf_ram_size: int,
+      apf_packet_format: int,
+  ):
+    self.apf_version_supported = apf_version_supported
+    self.apf_ram_size = apf_ram_size
+    self.apf_packet_format = apf_packet_format
+
+  def __str__(self):
+    """Returns a user-friendly string representation of the APF capabilities."""
+    return (
+        f"APF Version: {self.apf_version_supported}\n"
+        f"Ram Size: {self.apf_ram_size} bytes\n"
+        f"Packet Format: {self.apf_packet_format}"
+    )
+
+
+def get_apf_capabilities(
+    ad: android_device.AndroidDevice, iface_name: str
+) -> ApfCapabilities:
+  output = adb_utils.adb_shell(
+      ad, f"cmd network_stack apf {iface_name} capabilities"
+  )
+  try:
+    values = [int(value_str) for value_str in output.split(",")]
+  except ValueError:
+    return ApfCapabilities(0, 0, 0)  # Conversion to integer failed
+  return ApfCapabilities(values[0], values[1], values[2])
+
+
+def assume_apf_version_support_at_least(
+    ad: android_device.AndroidDevice, iface_name: str, expected_version: int
+) -> None:
+  caps = get_apf_capabilities(ad, iface_name)
+  asserts.skip_if(
+      caps.apf_version_supported < expected_version,
+      f"Supported apf version {caps.apf_version_supported} < expected version"
+      f" {expected_version}",
+  )