"""Tests for gce."""

from __future__ import print_function

import os

from chromite.lib import cros_test_lib
from chromite.lib import gce
from chromite.lib import osutils
from googleapiclient.errors import HttpError
from googleapiclient.http import HttpMockSequence
from oauth2client.client import GoogleCredentials


class GceTest(cros_test_lib.MockTempDirTestCase):
  """Unit tests for gce."""

  _PROJECT = 'foo-project'
  _ZONE = 'foo-zone'

  def setUp(self):
    self.json_key_file = os.path.join(self.tempdir, 'service_account.json')
    osutils.Touch(self.json_key_file)
    for cmd in ('from_stream', 'create_scoped'):
      self.PatchObject(GoogleCredentials, cmd, autospec=True)

  def testGetImage(self):
    """Tests that GetImage returns the correct image."""
    good_http = HttpMockSequence([
        ({'status': '200',}, '{"name": "foo-image"}',),
    ])

    # Assert that GetImage does not complain if an image is found.
    self.PatchObject(gce.GceContext, '_BuildRetriableRequest', autospec=True,
                     side_effect=self._MockOutBuildRetriableRequest(good_http))
    gce_context = gce.GceContext.ForServiceAccount(self._PROJECT, self._ZONE,
                                                   self.json_key_file)
    self.assertDictEqual(gce_context.GetImage('foo-image'),
                         dict(name='foo-image'))

  def testGetImageRaisesIfImageNotFound(self):
    """Tests that GetImage raies exception when image is not found."""
    bad_http = HttpMockSequence([
        ({'status': '404',}, 'Image not found.',),
    ])

    # Assert that GetImage raises if image is not found.
    self.PatchObject(gce.GceContext, '_BuildRetriableRequest', autospec=True,
                     side_effect=self._MockOutBuildRetriableRequest(bad_http))
    gce_context = gce.GceContext.ForServiceAccount(self._PROJECT, self._ZONE,
                                                   self.json_key_file)
    with self.assertRaises(gce.ResourceNotFoundError):
      gce_context.GetImage('not-exising-image')

  def testRetryOnServerErrorHttpRequest(self):
    """Tests that 500 erros are retried."""

    # Fake http sequence that does not return 200 until the third trial.
    mock_http = HttpMockSequence([
        ({'status': '502'}, 'Server error'),
        ({'status': '502'}, 'Server error'),
        ({'status': '200'}, '{"name":"foo-image"}'),
    ])

    self.PatchObject(gce.GceContext, '_BuildRetriableRequest', autospec=True,
                     side_effect=self._MockOutBuildRetriableRequest(mock_http))

    # Disable retry and expect the request to fail.
    gce.GceContext.RETRIES = 0
    gce_context = gce.GceContext.ForServiceAccount(self._PROJECT, self._ZONE,
                                                   self.json_key_file)
    with self.assertRaises(HttpError):
      gce_context.GetImage('foo-image')

    # Enable retry and expect the request to succeed.
    gce.GceContext.RETRIES = 2
    gce_context = gce.GceContext.ForServiceAccount(self._PROJECT, self._ZONE,
                                                   self.json_key_file)
    self.assertDictEqual(gce_context.GetImage('foo-image'),
                         dict(name='foo-image'))

  def _MockOutBuildRetriableRequest(self, mock_http):
    """Returns a mock closure of _BuildRetriableRequest."""
    # Replace the transport with mock_http.
    return (lambda self, num_retries, _http, _thread_safe, _credentials, *args,
            **kwargs:
            gce.RetryOnServerErrorHttpRequest(num_retries, mock_http, *args,
                                              **kwargs))


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