1#  Copyright (C) 2024 The Android Open Source Project
2#
3#  Licensed under the Apache License, Version 2.0 (the "License");
4#  you may not use this file except in compliance with the License.
5#  You may obtain a copy of the License at
6#
7#       http://www.apache.org/licenses/LICENSE-2.0
8#
9#  Unless required by applicable law or agreed to in writing, software
10#  distributed under the License is distributed on an "AS IS" BASIS,
11#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12#  See the License for the specific language governing permissions and
13#  limitations under the License.
14
15from mobly import asserts
16from mobly import base_test
17from net_tests_utils.host.python.assert_utils import (
18    UnexpectedBehaviorError, UnexpectedExceptionError, expect_with_retry, expect_throws
19)
20
21
22class TestAssertUtils(base_test.BaseTestClass):
23
24  def test_predicate_succeed(self):
25    """Test when the predicate becomes True within retries."""
26    call_count = 0
27
28    def predicate():
29      nonlocal call_count
30      call_count += 1
31      return call_count > 2  # True on the third call
32
33    expect_with_retry(predicate, max_retries=5, retry_interval_sec=0)
34    asserts.assert_equal(call_count, 3)  # Ensure it was called exactly 3 times
35
36  def test_predicate_failed(self):
37    """Test when the predicate never becomes True."""
38
39    with asserts.assert_raises(UnexpectedBehaviorError):
40      expect_with_retry(
41          predicate=lambda: False, max_retries=3, retry_interval_sec=0
42      )
43
44  def test_retry_action_not_called_succeed(self):
45    """Test that the retry_action is not called if the predicate returns true in the first try."""
46    retry_action_called = False
47
48    def retry_action():
49      nonlocal retry_action_called
50      retry_action_called = True
51
52    expect_with_retry(
53        predicate=lambda: True,
54        retry_action=retry_action,
55        max_retries=5,
56        retry_interval_sec=0,
57    )
58    asserts.assert_false(
59        retry_action_called, "retry_action called."
60    )  # Assert retry_action was NOT called
61
62  def test_retry_action_not_called_failed(self):
63    """Test that the retry_action is not called if the max_retries is reached."""
64    retry_action_called = False
65
66    def retry_action():
67      nonlocal retry_action_called
68      retry_action_called = True
69
70    with asserts.assert_raises(UnexpectedBehaviorError):
71      expect_with_retry(
72          predicate=lambda: False,
73          retry_action=retry_action,
74          max_retries=1,
75          retry_interval_sec=0,
76      )
77    asserts.assert_false(
78        retry_action_called, "retry_action called."
79    )  # Assert retry_action was NOT called
80
81  def test_retry_action_called(self):
82    """Test that the retry_action is executed when provided."""
83    retry_action_called = False
84
85    def retry_action():
86      nonlocal retry_action_called
87      retry_action_called = True
88
89    with asserts.assert_raises(UnexpectedBehaviorError):
90      expect_with_retry(
91          predicate=lambda: False,
92          retry_action=retry_action,
93          max_retries=2,
94          retry_interval_sec=0,
95      )
96    asserts.assert_true(retry_action_called, "retry_action not called.")
97
98  def test_expect_exception_throws(self):
99      def raise_unexpected_behavior_error():
100          raise UnexpectedBehaviorError()
101
102      expect_throws(raise_unexpected_behavior_error, UnexpectedBehaviorError)
103
104  def test_unexpect_exception_throws(self):
105      def raise_value_error():
106          raise ValueError()
107
108      with asserts.assert_raises(UnexpectedExceptionError):
109          expect_throws(raise_value_error, UnexpectedBehaviorError)
110
111  def test_no_exception_throws(self):
112      def raise_no_error():
113          return
114
115      expect_throws(raise_no_error, UnexpectedBehaviorError)