amemdiff_test.py

#
# Copyright (C) 2025 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
import argparse
import os
from unittest.mock import MagicMock, patch
import amemdiff
from amemdiff import _arg_key_value

class AmemdiffTest(unittest.TestCase):
    def test_arg_key_value_invalid(self):
        with self.assertRaises(argparse.ArgumentTypeError):
            _arg_key_value("key-value")

    def test_arg_key_value_valid(self):
        self.assertEqual(
            _arg_key_value("key/value"),
            {"key": "value"}
        )
        self.assertEqual(
            _arg_key_value("key1/value1,key2/value2"),
            {"key1": "value1", "key2": "value2"}
        )

    # pylint: disable=too-many-positional-arguments
    @patch('amemdiff.get_run_config')
    @patch('amemdiff.AndroidDevice')
    @patch('amemdiff.write_all_results')
    @patch('amemdiff.write_device_details_file')
    @patch('amemdiff.make_and_write_file')
    @patch('amemdiff.datetime')
    def test_main_writes_run_config_file(
        self,
        mock_datetime,
        mock_make_and_write_file,
        mock_write_device_details_file,
        mock_write_all_results,
        mock_android_device,
        mock_get_run_config
    ):
        """Test that main function writes the run_config.txt file."""
        # Arrange
        args = argparse.Namespace(
            devices={'d1': 's1', 'd2': 's2'},
            single_device_mode=False,
            output_directory='test_out',
            disable_device_configs=None,
            device_configs=None
        )
        run_config_str = "test run config"
        mock_run_config = MagicMock()
        mock_run_config.__str__.return_value = run_config_str
        mock_get_run_config.return_value = mock_run_config

        mock_device = MagicMock()
        mock_device.get_results.return_value = ["result1", "result2"]
        mock_android_device.return_value = mock_device

        mock_now = MagicMock()
        mock_now.strftime.return_value = '20251016_100000'
        mock_datetime.datetime.now.return_value = mock_now

        # Act
        amemdiff.main(args)

        # Assert
        expected_dir = os.path.join('test_out', '20251016_100000')
        mock_make_and_write_file.assert_called_with(
            expected_dir, 'run_config.txt', run_config_str
        )
        mock_write_all_results.assert_called_with(
            expected_dir, mock_run_config.database_connection
        )
        mock_write_device_details_file.assert_called_with(
            expected_dir, mock_run_config.database_connection
        )

    @patch('amemdiff.get_run_config')
    def test_main_exits_not_enough_devices(self, _mock_get_run_config):
        """Test that main exits if fewer than 2 devices are provided."""
        args = argparse.Namespace(
            devices={'d1': 's1'},
            single_device_mode=False,
            output_directory='test_out',
            disable_device_configs=None,
            device_configs=None
        )
        with self.assertRaises(SystemExit) as cm:
            amemdiff.main(args)

        self.assertEqual(
            cm.exception.code,
            (
                "1 device(s) found. Please specify two devices. E.g., "
                "-d=firstDevice/0.0.0.0:1234,secondDevice/0.0.0.0:1234."
            ),
        )

    @patch('amemdiff.get_run_config')
    def test_main_exits_too_many_devices(self, _mock_get_run_config):
        """Test that main exits if more than 2 devices are provided."""
        args = argparse.Namespace(
            devices={'d1': 's1', 'd2': 's2', 'd3': 's3'},
            single_device_mode=False,
            output_directory='test_out',
            disable_device_configs=None,
            device_configs=None
        )
        with self.assertRaises(SystemExit) as cm:
            amemdiff.main(args)

        self.assertEqual(
            cm.exception.code,
            (
                "3 device(s) found. Please specify two devices. E.g., "
                "-d=firstDevice/0.0.0.0:1234,secondDevice/0.0.0.0:1234."
            ),
        )

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