#!/usr/bin/env vpython3
# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import io
import json
import os
import sys
import unittest
from unittest import mock

import update_product_bundles

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),
                                             'test')))

import common
import compatible_utils


class TestUpdateProductBundles(unittest.TestCase):
  def setUp(self):
    # By default, test in attended mode.
    compatible_utils.force_running_attended()
    ffx_mock = mock.Mock()
    ffx_mock.returncode = 0
    self._ffx_patcher = mock.patch('common.run_ffx_command',
                                   return_value=ffx_mock)
    self._ffx_mock = self._ffx_patcher.start()
    self.addCleanup(self._ffx_mock.stop)

  def testConvertToProductBundleDefaultsUnknownImage(self):
    self.assertEqual(
        update_product_bundles.convert_to_products(['unknown-image']),
        ['unknown-image'])

  def testConvertToProductBundleRemovesReleaseSuffix(self):
    self.assertEqual(
        update_product_bundles.convert_to_products(
            ['smart_display_eng.astro-release']), ['smart_display_eng.astro'])

  def testConvertToProductBundleWarnsDeprecated(self):
    with self.assertLogs(level='WARNING') as logs:
      deprecated_images = [
          'qemu.arm64', 'qemu.x64', 'core.x64-dfv2-release',
          'workstation_eng.chromebook-x64-release'
      ]
      self.assertEqual(
          update_product_bundles.convert_to_products(deprecated_images), [
              'terminal.qemu-arm64', 'terminal.x64', 'core.x64-dfv2',
              'workstation_eng.chromebook-x64'
          ])
      for i, deprecated_image in enumerate(deprecated_images):
        self.assertIn(f'Image name {deprecated_image} has been deprecated',
                      logs.output[i])


  @mock.patch('common.run_ffx_command')
  def testRemoveRepositoriesRunsRemoveOnGivenRepos(self, ffx_mock):
    update_product_bundles.remove_repositories(['foo', 'bar', 'fizz', 'buzz'])

    ffx_mock.assert_has_calls([
        mock.call(cmd=('repository', 'remove', 'foo')),
        mock.call(cmd=('repository', 'remove', 'bar')),
        mock.call(cmd=('repository', 'remove', 'fizz')),
        mock.call(cmd=('repository', 'remove', 'buzz')),
    ])

  @mock.patch('os.path.exists')
  @mock.patch('os.path.abspath')
  def testGetRepositoriesPrunesReposThatDoNotExist(self, mock_abspath,
                                                   mock_exists):
    with mock.patch('common.SDK_ROOT', 'some/path'):
      self._ffx_mock.return_value.stdout = json.dumps([{
          "name": "terminal.x64",
          "spec": {
              "type": "pm",
              "path": "some/path/that/exists"
          }
      }, {
          "name": "workstation-eng.chromebook-x64",
          "spec": {
              "type": "pm",
              "path": "some/path/that/does/not/exist"
          }
      }])
      mock_exists.side_effect = [True, False]
      mock_abspath.side_effect = lambda x: x

      self.assertEqual(update_product_bundles.get_repositories(), [{
          "name": "terminal.x64",
          "spec": {
              "type": "pm",
              "path": "some/path/that/exists"
          }
      }])

      self._ffx_mock.assert_has_calls([
          mock.call(cmd=('--machine', 'json', 'repository', 'list'),
                    capture_output=True),
          mock.call(cmd=('repository', 'remove',
                         'workstation-eng.chromebook-x64'))
      ])

  @mock.patch('common.make_clean_directory')
  @mock.patch('update_product_bundles.running_unattended', return_value=True)
  # Disallow reading sdk_override.
  @mock.patch('os.path.isfile', return_value=False)
  @mock.patch('update_product_bundles.internal_hash', return_value='1.1.1')
  def testLookupAndDownloadWithAuth(self, *_):
    try:
      common.get_host_os()
    except:
      # Ignore unsupported platforms. common.get_host_os used in
      # update_product_bundles.main throws an unsupported exception.
      return
    auth_file = os.path.abspath(
        os.path.join(os.path.dirname(__file__), 'get_auth_token.py'))
    self._ffx_mock.return_value.stdout = json.dumps({
        "name": "core.x64",
        "product_version": "17.20240106.2.1",
        "transfer_manifest_url": "http://download-url"
    })

    with mock.patch(
        'sys.argv',
        ['update_product_bundles.py', 'terminal.x64', '--internal']):
      update_product_bundles.main()
    self._ffx_mock.assert_has_calls([
        mock.call(cmd=[
            '--machine', 'json', 'product', 'lookup', 'terminal.x64', '1.1.1',
            '--base-url', f'gs://fuchsia-sdk/development/1.1.1', '--auth',
            auth_file
        ],
                  capture_output=True),
        mock.call(cmd=[
            'product', 'download', 'http://download-url',
            os.path.join(common.INTERNAL_IMAGES_ROOT, 'terminal', 'x64'),
            '--auth', auth_file
        ])
    ])

  @mock.patch('common.make_clean_directory')
  @mock.patch('common.get_hash_from_sdk', return_value='2.2.2')
  @mock.patch('update_product_bundles.get_current_signature',
              return_value='2.2.2')
  @mock.patch('update_sdk.GetSDKOverrideGCSPath', return_value=None)
  def testIgnoreDownloadImagesWithSameHash(self, *_):
    try:
      common.get_host_os()
    except:
      # Ignore unsupported platforms. common.get_host_os used in
      # update_product_bundles.main throws an unsupported exception.
      return
    with mock.patch('sys.argv', ['update_product_bundles.py', 'terminal.x64']):
      update_product_bundles.main()
    self.assertFalse(self._ffx_mock.called)

  @mock.patch('common.make_clean_directory')
  @mock.patch('common.get_hash_from_sdk', return_value='2.2.2')
  @mock.patch('update_product_bundles.get_current_signature',
              return_value='0.0')
  @mock.patch('update_sdk.GetSDKOverrideGCSPath', return_value=None)
  def testDownloadImagesWithDifferentHash(self, *_):
    try:
      common.get_host_os()
    except:
      # Ignore unsupported platforms. common.get_host_os used in
      # update_product_bundles.main throws an unsupported exception.
      return
    self._ffx_mock.return_value.stdout = json.dumps({
        "name":
        "core.x64",
        "product_version":
        "17.20240106.2.1",
        "transfer_manifest_url":
        "http://download-url"
    })
    with mock.patch('sys.argv', ['update_product_bundles.py', 'terminal.x64']):
      update_product_bundles.main()
    self._ffx_mock.assert_has_calls([
        mock.call(cmd=[
            '--machine', 'json', 'product', 'lookup', 'terminal.x64', '2.2.2',
            '--base-url', 'gs://fuchsia/development/2.2.2'
        ],
                  capture_output=True),
        mock.call(cmd=[
            'product', 'download', 'http://download-url',
            os.path.join(common.IMAGES_ROOT, 'terminal', 'x64')
        ])
    ])

  @mock.patch('common.make_clean_directory')
  @mock.patch('update_sdk.GetSDKOverrideGCSPath',
              return_value='gs://my-bucket/sdk')
  @mock.patch('common.get_hash_from_sdk', return_value='2.2.2')
  def testSDKOverrideForSDKImages(self, *_):
    try:
      common.get_host_os()
    except:
      # Ignore unsupported platforms. common.get_host_os used in
      # update_product_bundles.main throws an unsupported exception.
      return
    self._ffx_mock.return_value.stdout = json.dumps({
        "name":
        "core.x64",
        "product_version":
        "17.20240106.2.1",
        "transfer_manifest_url":
        "http://download-url"
    })
    with mock.patch('sys.argv', ['update_product_bundles.py', 'terminal.x64']):
      update_product_bundles.main()
    self._ffx_mock.assert_has_calls([
        mock.call(cmd=[
            '--machine', 'json', 'product', 'lookup', 'terminal.x64', '2.2.2',
            '--base-url', 'gs://my-bucket'
        ],
                  capture_output=True),
        mock.call(cmd=[
            'product', 'download', 'http://download-url',
            os.path.join(common.IMAGES_ROOT, 'terminal', 'x64')
        ])
    ])

  @mock.patch('common.make_clean_directory')
  @mock.patch('update_product_bundles.internal_hash', return_value='1.1.1')
  @mock.patch('update_product_bundles.get_current_signature',
              return_value='1.1.1')
  def testIgnoreDownloadInternalImagesWithSameHash(self, *_):
    try:
      common.get_host_os()
    except:
      # Ignore unsupported platforms. common.get_host_os used in
      # update_product_bundles.main throws an unsupported exception.
      return
    with mock.patch(
        'sys.argv',
        ['update_product_bundles.py', 'terminal.x64', '--internal']):
      update_product_bundles.main()
    self.assertFalse(self._ffx_mock.called)

  @mock.patch('common.make_clean_directory')
  @mock.patch('update_product_bundles.get_current_signature',
              return_value='0.0')
  @mock.patch('update_product_bundles.internal_hash', return_value='1.1.1')
  def testDownloadInternalImagesWithDifferentHash(self, *_):
    try:
      common.get_host_os()
    except:
      # Ignore unsupported platforms. common.get_host_os used in
      # update_product_bundles.main throws an unsupported exception.
      return
    self._ffx_mock.return_value.stdout = json.dumps({
        "name":
        "core.x64",
        "product_version":
        "17.20240106.2.1",
        "transfer_manifest_url":
        "http://download-url"
    })
    with mock.patch(
        'sys.argv',
        ['update_product_bundles.py', 'terminal.x64', '--internal']):
      update_product_bundles.main()
    self._ffx_mock.assert_has_calls([
        mock.call(cmd=[
            '--machine', 'json', 'product', 'lookup', 'terminal.x64', '1.1.1',
            '--base-url', 'gs://fuchsia-sdk/development/1.1.1'
        ],
                  capture_output=True),
        mock.call(cmd=[
            'product', 'download', 'http://download-url',
            os.path.join(common.INTERNAL_IMAGES_ROOT, 'terminal', 'x64')
        ])
    ])

  @mock.patch('common.make_clean_directory')
  @mock.patch('update_sdk.GetSDKOverrideGCSPath',
              return_value='gs://my-bucket/sdk')
  @mock.patch('update_product_bundles.internal_hash', return_value='1.1.1')
  def testIgnoreSDKOverrideForInternalImages(self, *_):
    try:
      common.get_host_os()
    except:
      # Ignore unsupported platforms. common.get_host_os used in
      # update_product_bundles.main throws an unsupported exception.
      return
    self._ffx_mock.return_value.stdout = json.dumps({
        "name":
        "core.x64",
        "product_version":
        "17.20240106.2.1",
        "transfer_manifest_url":
        "http://download-url"
    })
    with mock.patch(
        'sys.argv',
        ['update_product_bundles.py', 'terminal.x64', '--internal']):
      update_product_bundles.main()
    self._ffx_mock.assert_has_calls([
        mock.call(cmd=[
            '--machine', 'json', 'product', 'lookup', 'terminal.x64', '1.1.1',
            '--base-url', 'gs://fuchsia-sdk/development/1.1.1'
        ],
                  capture_output=True),
        mock.call(cmd=[
            'product', 'download', 'http://download-url',
            os.path.join(common.INTERNAL_IMAGES_ROOT, 'terminal', 'x64')
        ])
    ])


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