# Copyright 2023 Google LLC
#
# 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.

"""Unittest Lab exercise to test implementation of "Synonym Dictionary"."""
import os
import subprocess
import sys
import unittest
from unittest import mock

from absl.testing import flagsaver
from mobly.controllers import android_device
from mobly.controllers.android_device_lib import adb
from mobly.snippet import errors as snippet_errors

from ui_automator import errors
from ui_automator import ui_automator
from ui_automator import version

_FAKE_MATTER_DEVICE_NAME = 'fake-matter-device-name'
_FAKE_GHA_ROOM = 'Office'
_FAKE_PAIRING_CODE = '34970112332'
_GOOGLE_HOME_APP = {
    'id': 'gha',
    'name': 'Google Home Application (GHA)',
    'minVersion': '3.1.18.1',
    'minMatter': '231456000',
    'minThread': '231456000',
}
_PYTHON_PATH = subprocess.check_output(['which', 'python']).decode('utf-8')
_PYTHON_BIN_PATH = _PYTHON_PATH.removesuffix('python')
_FAKE_VALID_SYS_ARGV_FOR_COMMISSIONING = [
    _PYTHON_BIN_PATH + 'ui-automator',
    '--commission',
    'm5stack,34970112332,Office',
]
_FAKE_VALID_SYS_ARGV_FOR_DECOMMISSIONING = [
    _PYTHON_BIN_PATH + 'ui-automator',
    '--decommission',
    'm5stack',
]
_FAKE_SYS_ARGV_FOR_COMMISSIONING_WITH_INVALID_LENGTH = [
    _PYTHON_BIN_PATH + 'ui-automator',
    '--commission',
    '',
]


class UIAutomatorTest(unittest.TestCase):

  def setUp(self):
    """This method will be run before each of the test methods in the class."""
    super().setUp()
    self.ui_automator = ui_automator.UIAutomator()
    self.mock_android_device = mock.patch.object(
        android_device, 'AndroidDevice'
    ).start()
    self.addCleanup(mock.patch.stopall)

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_load_device_raises_an_error_when_adb_not_installed(
      self, mock_get_all_instances
  ):
    mock_get_all_instances.side_effect = adb.AdbError(
        cmd='adb devices',
        stdout='fake_msg',
        stderr='adb command not found',
        ret_code=1,
    )
    with self.assertRaisesRegex(
        errors.AdbError,
        r'Please install adb and add it to your PATH environment variable\.',
    ):
      self.ui_automator.load_device()

  @mock.patch.object(
      android_device, 'get_all_instances', autospec=True, return_value=[]
  )
  def test_load_device_no_android_device_error(
      self, unused_mock_get_all_instances
  ):
    with self.assertRaises(errors.NoAndroidDeviceError):
      self.ui_automator.load_device()

  @mock.patch.object(
      android_device, 'get_all_instances', autospec=True, return_value=[1, 2]
  )
  def test_load_device_success(self, unused_mock_get_all_instances):
    with self.assertLogs() as cm:
      self.ui_automator.load_device()
    self.assertEqual(cm.output, ['INFO:root:connected device: [1]'])

  def test_load_snippet_without_load_device_raises_error(self):
    with self.assertRaises(errors.NoAndroidDeviceError):
      self.ui_automator.load_snippet()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_load_snippet_success(self, mock_get_all_instances):
    mock_get_all_instances.return_value = [self.mock_android_device]

    self.ui_automator.load_device()
    self.ui_automator.load_snippet()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_load_snippet_fails_with_server_start_pre_check_error(
      self, mock_get_all_instances
  ):
    self.mock_android_device.load_snippet.side_effect = (
        snippet_errors.ServerStartPreCheckError('fake_ad', 'fake_msg')
    )
    mock_get_all_instances.return_value = [self.mock_android_device]

    self.ui_automator.load_device()

    with self.assertRaises(errors.AndroidDeviceNotReadyError):
      self.ui_automator.load_snippet()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_load_snippet_fails_with_server_start_error(
      self, mock_get_all_instances
  ):
    self.mock_android_device.load_snippet.side_effect = (
        snippet_errors.ServerStartError('fake_ad', 'fake_msg')
    )
    mock_get_all_instances.return_value = [self.mock_android_device]

    self.ui_automator.load_device()

    with self.assertRaises(errors.AndroidDeviceNotReadyError):
      self.ui_automator.load_snippet()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_load_snippet_fails_with_protocol_error(self, mock_get_all_instances):
    self.mock_android_device.load_snippet.side_effect = (
        snippet_errors.ProtocolError('fake_ad', 'fake_msg')
    )
    mock_get_all_instances.return_value = [self.mock_android_device]

    self.ui_automator.load_device()

    with self.assertRaises(errors.AndroidDeviceNotReadyError):
      self.ui_automator.load_snippet()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_load_snippet_fails_with_snippet_error(self, mock_get_all_instances):
    self.mock_android_device.load_snippet.side_effect = (
        android_device.SnippetError('fake_device', 'fake_msg')
    )
    mock_get_all_instances.return_value = [self.mock_android_device]

    self.ui_automator.load_device()

    with self.assertLogs(level='DEBUG') as cm:
      self.ui_automator.load_snippet()
    self.assertEqual(
        cm.output,
        [
            "DEBUG:root:'fake_device'::Service<SnippetManagementService>"
            ' fake_msg'
        ],
    )

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  @mock.patch.object(
      os.path, 'dirname', autospec=True, return_value='/path/to/'
  )
  def test_load_snippet_installs_apk_when_apk_is_not_installed(
      self, mock_dirname, mock_get_all_instances
  ):
    self.mock_android_device.adb.shell.return_value = b'package:installed.apk\n'
    mock_get_all_instances.return_value = [self.mock_android_device]
    self.ui_automator.load_device()

    self.ui_automator.load_snippet()

    self.mock_android_device.adb.install.assert_called_once_with(
        ['-r', '-g', '/path/to/android/app/snippet-0.2.0-rc.0.apk']
    )
    self.mock_android_device.adb.uninstall.assert_not_called()
    mock_dirname.assert_called_once()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  @mock.patch.object(
      os.path, 'dirname', autospec=True, return_value='/path/to/'
  )
  def test_load_snippet_should_not_install_apk_when_correct_apk_installed(
      self, mock_dirname, mock_get_all_instances
  ):
    self.mock_android_device.adb.shell.side_effect = [
        b'package:com.chip.interop.moblysnippet\n',
        f'versionName={version.VERSION}'.encode('utf-8'),
    ]
    mock_get_all_instances.return_value = [self.mock_android_device]
    self.ui_automator.load_device()

    self.ui_automator.load_snippet()

    self.mock_android_device.adb.install.assert_not_called()
    self.mock_android_device.adb.uninstall.assert_not_called()
    mock_dirname.assert_not_called()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  @mock.patch.object(
      os.path, 'dirname', autospec=True, return_value='/path/to/'
  )
  def test_load_snippet_uninstalls_apk_before_installing_it_with_incorrect_apk(
      self, mock_dirname, mock_get_all_instances
  ):
    self.mock_android_device.adb.shell.side_effect = [
        b'package:com.chip.interop.moblysnippet\n',
        b'versionName=fake.version',
    ]
    mock_get_all_instances.return_value = [self.mock_android_device]
    self.ui_automator.load_device()

    self.ui_automator.load_snippet()

    self.mock_android_device.adb.uninstall.assert_called_with(
        'com.chip.interop.moblysnippet'
    )
    self.mock_android_device.adb.install.assert_called_once_with(
        ['-r', '-g', '/path/to/android/app/snippet-0.2.0-rc.0.apk']
    )
    mock_dirname.assert_called_once()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  @mock.patch.object(
      os.path, 'dirname', autospec=True, return_value='/path/to/'
  )
  def test_load_snippet_installs_apk_when_no_apk_installed(
      self, mock_dirname, mock_get_all_instances
  ):
    self.mock_android_device.adb.shell.side_effect = [
        b'package:installed.apk\n',
        b'Unable to find package\n',
    ]
    mock_get_all_instances.return_value = [self.mock_android_device]
    self.ui_automator.load_device()

    self.ui_automator.load_snippet()

    self.mock_android_device.adb.uninstall.assert_not_called()
    self.mock_android_device.adb.install.assert_called_once_with(
        ['-r', '-g', '/path/to/android/app/snippet-0.2.0-rc.0.apk']
    )
    mock_dirname.assert_called_once()

  @mock.patch.object(ui_automator.UIAutomator, 'load_device', autospec=True)
  @mock.patch.object(ui_automator.UIAutomator, 'load_snippet', autospec=True)
  def test_get_android_device_ready_raises_an_error_when_load_device_throws_an_error(
      self, mock_load_snippet, mock_load_device
  ):
    mock_load_device.side_effect = errors.NoAndroidDeviceError(
        'No Android device connected to the host computer.'
    )

    @ui_automator.get_android_device_ready
    def decorated_function(self):
      del self
      pass

    with self.assertRaisesRegex(
        errors.AndroidDeviceNotReadyError, r'decorated_function failed\.'
    ):
      decorated_function(self.ui_automator)

    mock_load_device.assert_called_once()
    mock_load_snippet.assert_not_called()

  @mock.patch.object(ui_automator.UIAutomator, 'load_device', autospec=True)
  @mock.patch.object(ui_automator.UIAutomator, 'load_snippet', autospec=True)
  def test_get_android_device_ready_raises_an_error_when_load_snippet_throws_an_error(
      self, mock_load_snippet, mock_load_device
  ):
    mock_load_snippet.side_effect = errors.AndroidDeviceNotReadyError(
        'Check device(fake-serial) has installed required apk.'
    )

    @ui_automator.get_android_device_ready
    def decorated_function(self):
      del self
      pass

    with self.assertRaisesRegex(
        errors.AndroidDeviceNotReadyError, r'decorated_function failed\.'
    ):
      decorated_function(self.ui_automator)

    mock_load_snippet.assert_called_once()
    mock_load_device.assert_called_once()

  @mock.patch.object(ui_automator.UIAutomator, 'load_device', autospec=True)
  @mock.patch.object(ui_automator.UIAutomator, 'load_snippet', autospec=True)
  def test_get_android_device_ready_raises_no_error_on_success(
      self, mock_load_snippet, mock_load_device
  ):
    @ui_automator.get_android_device_ready
    def decorated_function(self):
      del self
      pass

    decorated_function(self.ui_automator)

    mock_load_snippet.assert_called_once()
    mock_load_device.assert_called_once()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  @mock.patch.object(ui_automator.UIAutomator, 'load_snippet', autospec=True)
  def test_commission_device_raises_an_error_when_device_is_not_ready(
      self, mock_load_snippet, mock_get_all_instances
  ):
    mock_get_all_instances.return_value = [self.mock_android_device]
    error_msg = 'Check device(fake-serial) is ready.'
    mock_load_snippet.side_effect = errors.AndroidDeviceNotReadyError(error_msg)

    with self.assertRaisesRegex(
        errors.AndroidDeviceNotReadyError, 'commission_device failed.'
    ):
      self.ui_automator.commission_device(
          _FAKE_MATTER_DEVICE_NAME, _FAKE_PAIRING_CODE, _FAKE_GHA_ROOM
      )

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_commission_device_calls_a_method_in_snippet_with_desired_args(
      self, mock_get_all_instances
  ):
    mock_get_all_instances.return_value = [self.mock_android_device]
    expected_matter_device = {
        'id': '_fake_matter_device_name',
        'name': _FAKE_MATTER_DEVICE_NAME,
        'pairingCode': _FAKE_PAIRING_CODE,
        'roomName': _FAKE_GHA_ROOM,
    }

    self.ui_automator.commission_device(
        _FAKE_MATTER_DEVICE_NAME,
        _FAKE_PAIRING_CODE,
        _FAKE_GHA_ROOM,
    )

    self.mock_android_device.mbs.commissionDevice.assert_called_once_with(
        _GOOGLE_HOME_APP, expected_matter_device
    )

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_commission_device_fails_when_snippet_method_throws_an_error(
      self, mock_get_all_instances
  ):
    mock_get_all_instances.return_value = [self.mock_android_device]
    expected_error_message = (
        'Unable to continue automated commissioning process on'
        f' device({self.mock_android_device.device_info["serial"]}).'
    )
    self.mock_android_device.mbs.commissionDevice.side_effect = Exception(
        'Can not find next button in the page.'
    )

    with self.assertRaises(errors.MoblySnippetError) as exc:
      self.ui_automator.commission_device(
          _FAKE_MATTER_DEVICE_NAME,
          _FAKE_PAIRING_CODE,
          _FAKE_GHA_ROOM,
      )

    self.assertIn(expected_error_message, str(exc.exception))

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_commission_device_raises_an_error_when_device_name_exceeds_limit(
      self, mock_get_all_instances
  ):
    invalid_device_name = 'fakeDeviceNameWith25Chars'
    mock_get_all_instances.return_value = [self.mock_android_device]

    with self.assertRaisesRegex(
        ValueError,
        'Value of DeviceName is invalid. Device name should be no more than 24'
        ' characters.',
    ):
      self.ui_automator.commission_device(
          invalid_device_name,
          _FAKE_PAIRING_CODE,
          _FAKE_GHA_ROOM,
      )

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_commission_device_raises_an_error_when_pairing_code_is_invalid(
      self, mock_get_all_instances
  ):
    invalid_pairing_code = '123456789'
    mock_get_all_instances.return_value = [self.mock_android_device]

    with self.assertRaisesRegex(
        ValueError,
        'Value of PairingCode is invalid. Paring code should be 11-digit or'
        ' 21-digit numeric code.',
    ):
      self.ui_automator.commission_device(
          device_name=_FAKE_MATTER_DEVICE_NAME,
          pairing_code=invalid_pairing_code,
          gha_room=_FAKE_GHA_ROOM,
      )

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_commission_device_raises_an_error_when_gha_room_is_invalid(
      self, mock_get_all_instances
  ):
    invalid_gha_room = 'Attic 2'
    mock_get_all_instances.return_value = [self.mock_android_device]

    with self.assertRaisesRegex(
        ValueError,
        (
            'Value of GHARoom is invalid. Valid values for GHARoom: Attic|Back'
            ' door|Backyard|Basement|Bathroom|Bedroom|Den|Dining'
            ' Room|Entryway|Family Room|Front door|Front'
            ' Yard|Garage|Hallway|Kitchen|Living Room|Master'
            ' Bedroom|Office|Shed|Side door'
        ),
    ):
      self.ui_automator.commission_device(
          device_name=_FAKE_MATTER_DEVICE_NAME,
          pairing_code=_FAKE_PAIRING_CODE,
          gha_room=invalid_gha_room,
      )

  @flagsaver.flagsaver(
      (ui_automator._COMMISSION, ['m5stack', '34970112332', 'Office'])
  )
  @mock.patch.object(sys, 'exit', autospec=True)
  @mock.patch.object(
      ui_automator.UIAutomator, 'commission_device', autospec=True
  )
  def test_run_calls_commission_device_with_valid_arguments(
      self, mock_commission_device, mock_exit
  ):
    with mock.patch.object(sys, 'argv', _FAKE_VALID_SYS_ARGV_FOR_COMMISSIONING):
      ui_automator.run()

    mock_commission_device.assert_called_once_with(
        mock.ANY, 'm5stack', '34970112332', 'Office'
    )
    mock_exit.assert_called()

  @mock.patch.object(sys.stderr, 'write', autospec=True)
  @mock.patch.object(sys, 'exit', autospec=True)
  def test_commission_with_cmd_invalid_arg_should_stderr(
      self, mock_exit, mock_stderr_write
  ):
    with mock.patch.object(
        sys, 'argv', _FAKE_SYS_ARGV_FOR_COMMISSIONING_WITH_INVALID_LENGTH
    ):
      ui_automator.run()

    self.assertEqual(mock_stderr_write.call_count, 2)
    first_call_args = ''.join(mock_stderr_write.call_args_list[0][0])
    self.assertEqual(
        first_call_args,
        'FATAL Flags parsing error: flag --commission=[]: Use'
        ' --commission {DeviceName},{PairingCode},{GHARoom} to commission a'
        ' device to google fabric on GHA.\n',
    )
    second_call_args = ''.join(mock_stderr_write.call_args_list[1][0])
    self.assertEqual(
        second_call_args,
        'Pass --helpshort or --helpfull to see help on flags.\n',
    )
    self.mock_android_device.mbs.commissionDevice.assert_not_called()
    mock_exit.assert_called()

  @flagsaver.flagsaver((ui_automator._DECOMMISSION, 'm5stack'))
  @mock.patch.object(
      ui_automator.UIAutomator, 'commission_device', autospec=True
  )
  @mock.patch.object(sys, 'exit', autospec=True)
  @mock.patch.object(
      ui_automator.UIAutomator, 'decommission_device', autospec=True
  )
  def test_run_calls_decommission_device_with_valid_arguments(
      self, mock_decommission_device, mock_exit, mock_commission_device
  ):
    with mock.patch.object(
        sys, 'argv', _FAKE_VALID_SYS_ARGV_FOR_DECOMMISSIONING
    ):
      ui_automator.run()

    mock_commission_device.assert_not_called()
    mock_decommission_device.assert_called_once_with(mock.ANY, 'm5stack')
    mock_exit.assert_called()

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_decommission_device_raises_an_error_with_invalid_device_name(
      self, mock_get_all_instances
  ):
    mock_get_all_instances.return_value = [self.mock_android_device]
    expected_error_message = (
        'Value of DeviceName is invalid. Device name should be no more than 24'
        ' characters.'
    )
    with self.assertRaisesRegex(ValueError, expected_error_message):
      self.ui_automator.decommission_device('device_name_longer_than_24_chars')

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_decommission_device_fails_when_snippet_method_throws_an_error(
      self, mock_get_all_instances
  ):
    mock_get_all_instances.return_value = [self.mock_android_device]
    expected_error_message = (
        f'Unable to remove {_FAKE_MATTER_DEVICE_NAME} from GHA'
        f' on device({self.mock_android_device.device_info["serial"]}).'
    )
    self.mock_android_device.mbs.removeDevice.side_effect = Exception(
        'Can not remove the device.'
    )

    with self.assertRaises(errors.MoblySnippetError) as exc:
      self.ui_automator.decommission_device(_FAKE_MATTER_DEVICE_NAME)

    self.assertIn(expected_error_message, str(exc.exception))

  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
  def test_decommission_device_successfully_removes_a_device(
      self, mock_get_all_instances
  ):
    mock_get_all_instances.return_value = [self.mock_android_device]

    with self.assertLogs() as cm:
      self.ui_automator.decommission_device(_FAKE_MATTER_DEVICE_NAME)

    self.mock_android_device.mbs.removeDevice.assert_called_once_with(
        _FAKE_MATTER_DEVICE_NAME
    )
    self.assertEqual(
        cm.output[2], 'INFO:root:Successfully remove the device on GHA.'
    )


if __name__ == '__main__':
  unittest.main()