Add assert_utils for multidevice test

Test: atest NetworkStaticLibHostPythonTests
Bug: 335368434
Change-Id: I2a9ee97b84cb16791c0c5e3cb03120e67ffeb0af
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 4cf93a8..47e5d5e 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -267,6 +267,10 @@
     },
     {
       "name": "FrameworksNetTests"
+    },
+    // TODO: Move to presumit after meet SLO requirement.
+    {
+      "name": "NetworkStaticLibHostPythonTests"
     }
   ],
   "mainline-presubmit": [
diff --git a/staticlibs/tests/unit/Android.bp b/staticlibs/tests/unit/Android.bp
index fa466f8..5d537ff 100644
--- a/staticlibs/tests/unit/Android.bp
+++ b/staticlibs/tests/unit/Android.bp
@@ -57,3 +57,22 @@
     jarjar_rules: "jarjar-rules.txt",
     test_suites: ["device-tests"],
 }
+
+python_test_host {
+    name: "NetworkStaticLibHostPythonTests",
+    srcs: [
+        "host/python/*.py",
+    ],
+    main: "host/python/run_tests.py",
+    libs: [
+        "mobly",
+        "net-tests-utils-host-python-common",
+    ],
+    test_config: "host/python/test_config.xml",
+    test_suites: [
+        "general-tests",
+    ],
+    test_options: {
+        unit_test: true,
+    },
+}
diff --git a/staticlibs/tests/unit/host/python/assert_utils_test.py b/staticlibs/tests/unit/host/python/assert_utils_test.py
new file mode 100644
index 0000000..b970d14
--- /dev/null
+++ b/staticlibs/tests/unit/host/python/assert_utils_test.py
@@ -0,0 +1,89 @@
+#  Copyright (C) 2024 The Android Open Source Project
+#
+#  Licensed 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.
+
+import unittest
+from net_tests_utils.host.python.assert_utils import UnexpectedBehaviorError, expect_with_retry
+
+
+class TestAssertUtils(unittest.TestCase):
+
+  def test_predicate_succeed(self):
+    """Test when the predicate becomes True within retries."""
+    call_count = 0
+
+    def predicate():
+      nonlocal call_count
+      call_count += 1
+      return call_count > 2  # True on the third call
+
+    expect_with_retry(predicate, max_retries=5, retry_interval_sec=0)
+    self.assertEqual(call_count, 3)  # Ensure it was called exactly 3 times
+
+  def test_predicate_failed(self):
+    """Test when the predicate never becomes True."""
+
+    with self.assertRaises(UnexpectedBehaviorError):
+      expect_with_retry(
+          predicate=lambda: False, max_retries=3, retry_interval_sec=0
+      )
+
+  def test_retry_action_not_called_succeed(self):
+    """Test that the retry_action is not called if the predicate returns true in the first try."""
+    retry_action_called = False
+
+    def retry_action():
+      nonlocal retry_action_called
+      retry_action_called = True
+
+    expect_with_retry(
+        predicate=lambda: True,
+        retry_action=retry_action,
+        max_retries=5,
+        retry_interval_sec=0,
+    )
+    self.assertFalse(retry_action_called)  # Assert retry_action was NOT called
+
+  def test_retry_action_not_called_failed(self):
+    """Test that the retry_action is not called if the max_retries is reached."""
+    retry_action_called = False
+
+    def retry_action():
+      nonlocal retry_action_called
+      retry_action_called = True
+
+    with self.assertRaises(UnexpectedBehaviorError):
+      expect_with_retry(
+          predicate=lambda: False,
+          retry_action=retry_action,
+          max_retries=1,
+          retry_interval_sec=0,
+      )
+    self.assertFalse(retry_action_called)  # Assert retry_action was NOT called
+
+  def test_retry_action_called(self):
+    """Test that the retry_action is executed when provided."""
+    retry_action_called = False
+
+    def retry_action():
+      nonlocal retry_action_called
+      retry_action_called = True
+
+    with self.assertRaises(UnexpectedBehaviorError):
+      expect_with_retry(
+          predicate=lambda: False,
+          retry_action=retry_action,
+          max_retries=2,
+          retry_interval_sec=0,
+      )
+    self.assertTrue(retry_action_called)
diff --git a/staticlibs/tests/unit/host/python/run_tests.py b/staticlibs/tests/unit/host/python/run_tests.py
new file mode 100644
index 0000000..8059568
--- /dev/null
+++ b/staticlibs/tests/unit/host/python/run_tests.py
@@ -0,0 +1,25 @@
+#  Copyright (C) 2024 The Android Open Source Project
+#
+#  Licensed 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.
+
+"""Main entrypoint for all of unittest."""
+
+import unittest
+
+# Import all unittest classes here, so it can be discovered by unittest module.
+# TODO: make the tests can be executed without manually import classes.
+from host.python.assert_utils_test import TestAssertUtils
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/staticlibs/tests/unit/host/python/test_config.xml b/staticlibs/tests/unit/host/python/test_config.xml
new file mode 100644
index 0000000..26ee9e2
--- /dev/null
+++ b/staticlibs/tests/unit/host/python/test_config.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed 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.
+-->
+<configuration description="Config for NetworkStaticLibHostPythonTests">
+    <target_preparer class="com.android.tradefed.targetprep.RootTargetPreparer" />
+    <test class="com.android.tradefed.testtype.python.PythonBinaryHostTest">
+        <option name="par-file-name" value="NetworkStaticLibHostPythonTests" />
+        <option name="test-timeout" value="3m" />
+    </test>
+</configuration>
diff --git a/staticlibs/testutils/host/python/assert_utils.py b/staticlibs/testutils/host/python/assert_utils.py
new file mode 100644
index 0000000..da1bb9e
--- /dev/null
+++ b/staticlibs/testutils/host/python/assert_utils.py
@@ -0,0 +1,43 @@
+#  Copyright (C) 2024 The Android Open Source Project
+#
+#  Licensed 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.
+
+import time
+from typing import Callable
+
+
+class UnexpectedBehaviorError(Exception):
+  """Raised when there is an unexpected behavior during applying a procedure."""
+
+
+def expect_with_retry(
+    predicate: Callable[[], bool],
+    retry_action: Callable[[], None] = None,
+    max_retries: int = 10,
+    retry_interval_sec: int = 1,
+) -> None:
+  """Executes a predicate and retries if it doesn't return True."""
+
+  for retry in range(max_retries):
+    if predicate():
+      return None
+    else:
+      if retry == max_retries - 1:
+        break
+      if retry_action:
+        retry_action()
+      time.sleep(retry_interval_sec)
+
+  raise UnexpectedBehaviorError(
+      "Predicate didn't become true after " + str(max_retries) + " retries."
+  )