tests/wlan_policy/RegulatoryRecoveryTest.py

#!/usr/bin/env python3
#
# Copyright 2022 The Fuchsia Authors
#
# 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.

from honeydew.typing.wlan import CountryCode, WlanClientState
from mobly import signals, test_runner

from antlion.controllers.ap_lib.hostapd_security import FuchsiaSecurityType
from antlion.controllers.fuchsia_device import FuchsiaDevice
from antlion.controllers.fuchsia_lib.lib_controllers.wlan_policy_controller import (
    WlanPolicyControllerError,
)
from antlion.controllers.fuchsia_lib.wlan_ap_policy_lib import (
    ConnectivityMode,
    OperatingBand,
)
from antlion.test_utils.wifi import base_test


class RegulatoryRecoveryTest(base_test.WifiBaseTest):
    """Tests the policy layer's response to setting country code.

    Test Bed Requirements:
    * One Fuchsia device that is capable of operating as a WLAN client and AP.

    Example Config:
    "regulatory_recovery_test_params": {
        "country_code": "US"
    }

    If no configuration information is provided, the test will default to
    toggling between WW and US.
    """

    def setup_class(self) -> None:
        super().setup_class()
        if len(self.fuchsia_devices) < 1:
            raise EnvironmentError("No Fuchsia devices found.")

        self.config_test_params = self.user_params.get(
            "regulatory_recovery_test_params", {}
        )
        self.country_code = self.config_test_params.get("country_code", "US")
        self.negative_test = self.config_test_params.get("negative_test", False)

        for fd in self.fuchsia_devices:
            fd.configure_wlan(association_mechanism="policy")

    def teardown_class(self) -> None:
        if not self.negative_test:
            for fd in self.fuchsia_devices:
                fd.wlan_controller.set_country_code(self.country_code)

        super().teardown_class()

    def setup_test(self) -> None:
        """Set PHYs to world-wide mode and disable AP and client connections."""
        for fd in self.fuchsia_devices:
            fd.wlan_controller.set_country_code(CountryCode.WORLDWIDE)
            fd.sl4f.wlan_ap_policy_lib.wlanStopAllAccessPoint()

    def _set_country_code_check(self, fd: FuchsiaDevice) -> None:
        """Set the country code and check if successful.

        Args:
            fd: Fuchsia device to set country code on.

        Raises:
            EnvironmentError on failure to set country code or success setting country
                code when it should be a failure case.
        """
        try:
            fd.wlan_controller.set_country_code(self.country_code)
        except EnvironmentError as e:
            if self.negative_test:
                # In the negative case, setting the country code for an
                # invalid country should fail.
                pass
            else:
                # If this is not a negative test case, re-raise the
                # exception.
                raise e
        else:
            # The negative test case should have failed to set the country
            # code and the positive test case should succeed.
            if self.negative_test:
                raise EnvironmentError("Setting invalid country code succeeded.")
            else:
                pass

    def test_interfaces_not_recreated_when_initially_disabled(self) -> None:
        """Test after applying new region no new interfaces are automatically recreated.

        We start with client connections and access points disabled. There should be no
        state change after applying a new regulatory region.

        Raises:
            TestFailure if client or AP are in unexpected state.
        """
        for fd in self.fuchsia_devices:
            fd.sl4f.wlan_policy_lib.stop_client_connections()
            fd.wlan_policy_controller.wait_for_client_state(
                WlanClientState.CONNECTIONS_DISABLED
            )

            self._set_country_code_check(fd)

            # Verify that the client is still stopped.
            try:
                fd.wlan_policy_controller.wait_for_client_state(
                    WlanClientState.CONNECTIONS_DISABLED
                )
            except WlanPolicyControllerError:
                raise signals.TestFailure("Client policy layer is in unexpected state")

            # Verify that the AP is still stopped.
            fd.sl4f.wlan_ap_policy_lib.wlanSetNewListener()
            ap_state = fd.sl4f.wlan_ap_policy_lib.wlanGetUpdate()
            if ap_state["error"]:
                raise signals.TestFailure(
                    f"error querying AP state: {ap_state['error']}"
                )

            ap_updates = ap_state["result"]
            if ap_updates:
                raise signals.TestFailure(f"AP in unexpected state: {ap_updates}")

    def test_interfaces_recreated_when_initially_enabled(self) -> None:
        """Test after applying new region interfaces are automatically recreated.

        After enabling client connections and access points we check that all interfaces
        are recreated.

        Raises:
            TestFailure if client or AP are in unexpected state.
        """
        test_ssid = "test_ssid"
        security_type = FuchsiaSecurityType.NONE
        for fd in self.fuchsia_devices:
            # Start client connections and start an AP before setting the country code.
            fd.sl4f.wlan_policy_lib.start_client_connections()
            fd.wlan_policy_controller.wait_for_client_state(
                WlanClientState.CONNECTIONS_ENABLED
            )
            fd.sl4f.wlan_ap_policy_lib.wlanStartAccessPoint(
                test_ssid,
                security_type,
                None,
                ConnectivityMode.LOCAL_ONLY,
                OperatingBand.ANY,
            )

            # Set the country code.
            self._set_country_code_check(fd)

            # Verify that the client connections are enabled.
            try:
                fd.wlan_policy_controller.wait_for_client_state(
                    WlanClientState.CONNECTIONS_ENABLED
                )
            except WlanPolicyControllerError:
                raise signals.TestFailure("Client policy layer is in unexpected state")

            # Verify that the AP is brought up again.
            fd.sl4f.wlan_ap_policy_lib.wlanSetNewListener()
            ap_state = fd.sl4f.wlan_ap_policy_lib.wlanGetUpdate()
            if ap_state["error"]:
                raise signals.TestFailure(
                    f"error querying AP state: {ap_state['error']}"
                )

            ap_updates = ap_state["result"]
            if len(ap_updates) != 1:
                raise signals.TestFailure(f"No APs are running: {ap_updates}")
            else:
                if (
                    ap_updates[0]["id"]["ssid"] != test_ssid
                    or ap_updates[0]["id"]["type_"].lower() != security_type
                ):
                    raise signals.TestFailure(
                        f"AP in unexpected state: {ap_updates[0]}"
                    )


if __name__ == "__main__":
    test_runner.main()