#!/usr/bin/env python3
"""Test for the ee.data module."""

from unittest import mock

import httplib2
import requests

import ee
from ee import _cloud_api_utils
from ee import apitestcase
from ee import featurecollection
from ee import image
import unittest


class DataTest(unittest.TestCase):

  def testListOperations(self):
    mock_http = mock.MagicMock(httplib2.Http)
    # Return in three groups.
    mock_http.request.side_effect = [
        (httplib2.Response({'status': 200}),
         b'{"operations": [{"name": "name1"}], "nextPageToken": "t1"}'),
        (httplib2.Response({'status': 200}),
         b'{"operations": [{"name": "name2"}], "nextPageToken": "t2"}'),
        (httplib2.Response({'status': 200}),
         b'{"operations": [{"name": "name3"}]}'),
    ]
    with apitestcase.UsingCloudApi(mock_http=mock_http):
      self.assertEqual([{
          'name': 'name1'
      }, {
          'name': 'name2'
      }, {
          'name': 'name3'
      }], ee.data.listOperations())

  def testListOperationsEmptyList(self):
    # Empty lists don't appear at all in the result.
    mock_http = mock.MagicMock(httplib2.Http)
    mock_http.request.return_value = (httplib2.Response({'status': 200}), b'{}')
    with apitestcase.UsingCloudApi(mock_http=mock_http):
      self.assertEqual([], ee.data.listOperations())

  def testCreateAsset(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {
          'type': 'FOLDER',
          'name': 'projects/earthengine-legacy/assets/users/foo/xyz1234',
          'id': 'users/foo/xyz1234',
      }
      cloud_api_resource.projects().assets().create.execute.return_value = (
          mock_result
      )
      ee.data.createAsset({'type': 'FOLDER'}, 'users/foo/xyz123')
      mock_create_asset = cloud_api_resource.projects().assets().create
      mock_create_asset.assert_called_once()
      parent = mock_create_asset.call_args.kwargs['parent']
      self.assertEqual(parent, 'projects/earthengine-legacy')
      asset_id = mock_create_asset.call_args.kwargs['assetId']
      self.assertEqual(asset_id, 'users/foo/xyz123')
      asset = mock_create_asset.call_args.kwargs['body']
      self.assertEqual(asset, {'type': 'FOLDER'})

  def testCreateAssetWithV1AlphaParams(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {
          'type': 'IMAGE',
          'name': 'projects/earthengine-legacy/assets/users/foo/xyz1234',
          'id': 'users/foo/xyz1234',
          'properties': {
              'title': 'My Test Asset',
              'description': 'original description',
              'myProperty': 1,
          },
          'cloudStorageLocation': {'uris': ['gs://my-bucket/path']},
          'tilestoreLocation': {'sources': []},
      }
      cloud_api_resource.projects().assets().create.execute.return_value = (
          mock_result
      )
      test_properties = {
          'myProperty': 1,
          'description': 'original description',
      }
      ee.data.createAsset(
          {
              'type': 'IMAGE',
              'gcs_location': {'uris': ['gs://my-bucket/path']},
              'tilestore_entry': {'sources': []},
              'title': 'My Test Asset',
              'description': 'new description',
              'properties': test_properties,
          },
          'users/foo/xyz123',
      )
      mock_create_asset = cloud_api_resource.projects().assets().create
      mock_create_asset.assert_called_once()
      parent = mock_create_asset.call_args.kwargs['parent']
      self.assertEqual(parent, 'projects/earthengine-legacy')
      asset_id = mock_create_asset.call_args.kwargs['assetId']
      self.assertEqual(asset_id, 'users/foo/xyz123')
      asset = mock_create_asset.call_args.kwargs['body']
      self.assertEqual(
          asset['properties'],
          {
              'title': 'My Test Asset',
              'description': 'original description',
              'myProperty': 1,
          },
      )
      self.assertEqual(test_properties, {
          'myProperty': 1,
          'description': 'original description',
      })
      self.assertEqual(
          asset['cloud_storage_location'],
          {'uris': ['gs://my-bucket/path']},
      )
      self.assertEqual(
          asset['tilestore_location'],
          {'sources': []},
      )

  def testSetAssetProperties(self):
    mock_http = mock.MagicMock(httplib2.Http)
    with apitestcase.UsingCloudApi(mock_http=mock_http), mock.patch.object(
        ee.data, 'updateAsset', autospec=True) as mock_update_asset:
      ee.data.setAssetProperties('foo', {
          'mYPropErTy': 'Value',
          'system:time_start': 1
      })
      asset_id = mock_update_asset.call_args[0][0]
      self.assertEqual(asset_id, 'foo')
      asset = mock_update_asset.call_args[0][1]
      self.assertEqual(asset['properties'], {
          'mYPropErTy': 'Value',
          'system:time_start': 1
      })
      update_mask = mock_update_asset.call_args[0][2]
      self.assertSetEqual(
          set(update_mask),
          set(['properties.\"mYPropErTy\"',
               'properties.\"system:time_start\"']))

  def testListAssets(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'assets': [{'path': 'id1', 'type': 'type1'}]}
      cloud_api_resource.projects().assets().listAssets(
      ).execute.return_value = mock_result
      cloud_api_resource.projects().assets().listAssets_next.return_value = None
      actual_result = ee.data.listAssets({'p': 'q'})
      cloud_api_resource.projects().assets().listAssets(
      ).execute.assert_called_once()
      self.assertEqual(mock_result, actual_result)

  def testListImages(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'assets': [{'path': 'id1', 'type': 'type1'}]}
      cloud_api_resource.projects().assets().listAssets(
      ).execute.return_value = mock_result
      cloud_api_resource.projects().assets().listAssets_next.return_value = None
      actual_result = ee.data.listImages({'p': 'q'})
      cloud_api_resource.projects().assets().listAssets(
      ).execute.assert_called_once()
      self.assertEqual({'images': [{
          'path': 'id1',
          'type': 'type1'
      }]}, actual_result)

  def testListBuckets(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'assets': [{'name': 'id1', 'type': 'FOLDER'}]}
      cloud_api_resource.projects().listAssets(
      ).execute.return_value = mock_result
      actual_result = ee.data.listBuckets()
    cloud_api_resource.projects().listAssets().execute.assert_called_once()
    self.assertEqual(mock_result, actual_result)

  def testSimpleGetListViaCloudApi(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'assets': [{'name': 'id1', 'type': 'IMAGE_COLLECTION'}]}
      cloud_api_resource.projects().assets().listAssets(
      ).execute.return_value = mock_result
      actual_result = ee.data.getList({'id': 'glam', 'num': 3})
      expected_params = {
          'parent': 'projects/earthengine-public/assets/glam',
          'pageSize': 3,
          'view': 'BASIC',
      }
      expected_result = [{'id': 'id1', 'type': 'ImageCollection'}]
      cloud_api_resource.projects().assets().listAssets.assert_called_with(
          **expected_params)
      self.assertEqual(expected_result, actual_result)

  def testGetListAssetRootViaCloudApi(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'assets': [{'name': 'id1', 'type': 'IMAGE_COLLECTION'}]}
      cloud_api_resource.projects().listAssets(
      ).execute.return_value = mock_result
      actual_result = ee.data.getList({
          'id': 'projects/my-project/assets/',
          'num': 3
      })
      expected_params = {
          'parent': 'projects/my-project',
          'pageSize': 3,
          'view': 'BASIC'
      }
      expected_result = [{'id': 'id1', 'type': 'ImageCollection'}]
      cloud_api_resource.projects().listAssets.assert_called_with(
          **expected_params)
      self.assertEqual(expected_result, actual_result)

  def testGetListAssetRootViaCloudApiNoSlash(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'assets': [{'name': 'id1', 'type': 'IMAGE_COLLECTION'}]}
      cloud_api_resource.projects().listAssets(
      ).execute.return_value = mock_result
      actual_result = ee.data.getList({
          'id': 'projects/my-project/assets',
          'num': 3
      })
      expected_params = {
          'parent': 'projects/my-project',
          'pageSize': 3,
          'view': 'BASIC'
      }
      expected_result = [{'id': 'id1', 'type': 'ImageCollection'}]
      cloud_api_resource.projects().listAssets.assert_called_with(
          **expected_params)
      self.assertEqual(expected_result, actual_result)

  def testComplexGetListViaCloudApi(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {
          'assets': [{
              'name': 'id1',
              'type': 'IMAGE',
              'size_bytes': 1234
          }]
      }
      cloud_api_resource.projects().assets().listAssets(
      ).execute.return_value = mock_result
      actual_result = ee.data.getList({
          'id': 'glam',
          'num': 3,
          'starttime': 3612345,
          'filter': 'foo'
      })
      expected_params = {
          'parent': 'projects/earthengine-public/assets/glam',
          'pageSize': 3,
          'view': 'BASIC',
          'filter': 'foo AND startTime >= "1970-01-01T01:00:12.345000Z"'
      }
      expected_result = [{'id': 'id1', 'type': 'Image'}]
      cloud_api_resource.projects().assets().listAssets.assert_called_with(
          **expected_params)
      self.assertEqual(expected_result, actual_result)

  def testGetMapId(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {
          'name': 'projects/earthengine-legacy/maps/DOCID',
      }
      cloud_api_resource.projects().maps().create(
      ).execute.return_value = mock_result
      actual_result = ee.data.getMapId({
          'image': image.Image('my-image'),
      })
      cloud_api_resource.projects().maps().create().execute.assert_called_once()
      self.assertEqual('projects/earthengine-legacy/maps/DOCID',
                       actual_result['mapid'])
      self.assertEqual('', actual_result['token'])
      self.assertIsInstance(actual_result['tile_fetcher'], ee.data.TileFetcher)

  def testGetMapId_withWorkloadTag(self):
    with ee.data.workloadTagContext('mapid-tag'):
      cloud_api_resource = mock.MagicMock()
      with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
        mock_result = {
            'name': 'projects/earthengine-legacy/maps/DOCID',
        }
        cloud_api_resource.projects().maps().create(
        ).execute.return_value = mock_result
        ee.data.getMapId({
            'image': image.Image('my-image'),
        })
        self.assertEqual(
            'mapid-tag',
            cloud_api_resource.projects().maps().create.call_args_list[1]
            .kwargs['workloadTag'])

  # The Cloud API context manager does not mock getAlgorithms, so it's done
  # separately here.
  @mock.patch.object(
      ee.data,
      'getAlgorithms',
      return_value=apitestcase.GetAlgorithms(),
      autospec=True)
  def testGetDownloadId(self, _):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'name': 'projects/earthengine-legacy/thumbnails/DOCID'}
      cloud_api_resource.projects().thumbnails().create(
      ).execute.return_value = mock_result
      actual_result = ee.data.getDownloadId({
          'image': image.Image('my-image'),
          'name': 'dummy'
      })
      cloud_api_resource.projects().thumbnails().create(
      ).execute.assert_called_once()
      self.assertEqual(
          {
              'docid': 'projects/earthengine-legacy/thumbnails/DOCID',
              'token': ''
          }, actual_result)

  def testGetDownloadId_withWorkloadTag(self):
    with ee.data.workloadTagContext('downloadid-tag'):
      cloud_api_resource = mock.MagicMock()
      with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
        mock_result = {'name': 'projects/earthengine-legacy/thumbnails/DOCID'}
        cloud_api_resource.projects().thumbnails().create(
        ).execute.return_value = mock_result
        ee.data.getDownloadId({
            'image': image.Image('my-image'),
            'name': 'dummy'
        })
        self.assertEqual(
            'downloadid-tag',
            cloud_api_resource.projects().thumbnails().create.call_args
            .kwargs['workloadTag'])

  def testGetDownloadId_withBandList(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'name': 'projects/earthengine-legacy/thumbnails/DOCID'}
      cloud_api_resource.projects().thumbnails().create(
      ).execute.return_value = mock_result
      actual_result = ee.data.getDownloadId({
          'image': image.Image('my-image'),
          'name': 'dummy',
          'bands': ['B1', 'B2', 'B3']
      })
      cloud_api_resource.projects().thumbnails().create(
      ).execute.assert_called_once()
      self.assertEqual(
          {
              'docid': 'projects/earthengine-legacy/thumbnails/DOCID',
              'token': ''
          }, actual_result)

  def testGetDownloadId_withImageID(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      with self.assertRaisesRegex(ee.ee_exception.EEException,
                                  '^Image ID string is not supported.'):
        ee.data.getDownloadId({'id': 'my-image', 'name': 'dummy'})

  def testGetDownloadId_withSerializedImage(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      with self.assertRaisesRegex(ee.ee_exception.EEException,
                                  '^Image as JSON string not supported.'):
        ee.data.getDownloadId({
            'image': image.Image('my-image').serialize(),
            'name': 'dummy'
        })

  def testGetThumbId(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'name': 'projects/earthengine-legacy/thumbnails/DOCID'}
      cloud_api_resource.projects().thumbnails().create(
      ).execute.return_value = mock_result
      actual_result = ee.data.getThumbId({
          'image': image.Image('my-image'),
          'name': 'dummy'
      })
      cloud_api_resource.projects().thumbnails().create(
      ).execute.assert_called_once()
      self.assertEqual(
          {
              'thumbid': 'projects/earthengine-legacy/thumbnails/DOCID',
              'token': ''
          }, actual_result)

  def testGetThumbId_withWorkloadTag(self):
    with ee.data.workloadTagContext('thumbid-tag'):
      cloud_api_resource = mock.MagicMock()
      with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
        mock_result = {'name': 'projects/earthengine-legacy/thumbnails/DOCID'}
        cloud_api_resource.projects().thumbnails().create(
        ).execute.return_value = mock_result
        ee.data.getThumbId({'image': image.Image('my-image'), 'name': 'dummy'})
        self.assertEqual(
            'thumbid-tag',
            cloud_api_resource.projects().thumbnails().create.call_args
            .kwargs['workloadTag'])

  def testGetTableDownloadId(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {'name': 'projects/earthengine-legacy/table/DOCID'}
      cloud_api_resource.projects().tables().create(
      ).execute.return_value = mock_result
      actual_result = ee.data.getTableDownloadId({
          'table': featurecollection.FeatureCollection('my-fc'),
          'filename': 'dummy'
      })
      cloud_api_resource.projects().tables().create(
      ).execute.assert_called_once()
      self.assertEqual(
          {
              'docid': 'projects/earthengine-legacy/table/DOCID',
              'token': ''
          }, actual_result)

  def testGetTableDownloadId_withWorkloadTag(self):
    with ee.data.workloadTagContext('tableid-tag'):
      cloud_api_resource = mock.MagicMock()
      with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
        mock_result = {'name': 'projects/earthengine-legacy/thumbnails/DOCID'}
        cloud_api_resource.projects().tables().create(
        ).execute.return_value = mock_result
        ee.data.getTableDownloadId({
            'table': featurecollection.FeatureCollection('my-fc'),
            'filename': 'dummy'
        })
        self.assertEqual(
            'tableid-tag',
            cloud_api_resource.projects().tables().create.call_args
            .kwargs['workloadTag'])

  def testCloudProfilingEnabled(self):
    seen = []

    def ProfileHook(profile_id):
      seen.append(profile_id)

    with ee.data.profiling(ProfileHook):
      with apitestcase.UsingCloudApi(), DoCloudProfileStubHttp(self, True):
        ee.data.listImages({'parent': 'projects/earthengine-public/assets/q'})
    self.assertEqual(['someProfileId'], seen)

  def testCloudProfilingDisabled(self):
    with apitestcase.UsingCloudApi(), DoCloudProfileStubHttp(self, False):
      ee.data.listImages({'parent': 'projects/earthengine-public/assets/q'})

  def testCloudErrorTranslation(self):
    mock_http = mock.MagicMock(httplib2.Http)
    mock_http.request.return_value = (httplib2.Response({'status': 400}),
                                      b'{"error": {"message": "errorly"} }')
    with apitestcase.UsingCloudApi(mock_http=mock_http):
      with self.assertRaisesRegex(ee.ee_exception.EEException, '^errorly$'):
        ee.data.listImages({'parent': 'projects/earthengine-public/assets/q'})

  def testListFeatures(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_result = {
          'type':
              'FeatureCollection',
          'features': [{
              'type': 'Feature',
              'properties': {
                  'baz': 'qux',
                  'foo': 'bar',
                  'system:index': '0'
              }
          }]
      }
      cloud_api_resource.projects().assets().listFeatures(
      ).execute.return_value = mock_result
      actual_result = ee.data.listFeatures({
          'assetId':
              'users/userfoo/foobar',
          'region':
              '{\"type\":\"Polygon\",\"coordinates\":[[[-96,42],[-95,42],[-95,43],[-96,43],[-96,42]]]}'
      })
      cloud_api_resource.projects().assets().listFeatures(
      ).execute.assert_called_once()
      self.assertEqual(mock_result, actual_result)

  @mock.patch.object(ee.data, '_tile_base_url', new='base_url')
  def testGetFeatureViewTilesKey(self):
    cloud_api_resource = mock.MagicMock()
    with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
      mock_name = 'projects/projectfoo/featureView/tiles-key-foo'
      mock_result = {'name': mock_name}
      cloud_api_resource.projects().featureView().create(
      ).execute.return_value = mock_result
      actual_result = ee.data.getFeatureViewTilesKey({
          'assetId': 'projects/projectfoo/assets/assetbar',
      })
      cloud_api_resource.projects().featureView().create(
      ).execute.assert_called_once()
      expected_keys = [
          'token',
          'formatTileUrl',
      ]
      self.assertEqual(expected_keys, list(actual_result.keys()))
      self.assertEqual('tiles-key-foo', actual_result['token'])
      self.assertEqual(
          f'base_url/{_cloud_api_utils.VERSION}/{mock_name}/tiles/7/5/6',
          actual_result['formatTileUrl'](5, 6, 7))

  def testWorkloadTag(self):
    self.assertEqual('', ee.data.getWorkloadTag())
    ee.data.setDefaultWorkloadTag(None)
    self.assertEqual('', ee.data.getWorkloadTag())
    ee.data.setDefaultWorkloadTag('')
    self.assertEqual('', ee.data.getWorkloadTag())
    ee.data.setDefaultWorkloadTag(0)
    self.assertEqual('0', ee.data.getWorkloadTag())
    ee.data.setDefaultWorkloadTag(123)
    self.assertEqual('123', ee.data.getWorkloadTag())

    with self.assertRaisesRegex(ValueError, 'Invalid tag'):
      ee.data.setDefaultWorkloadTag('inv@lid')

    with self.assertRaisesRegex(ValueError, 'Invalid tag'):
      ee.data.setDefaultWorkloadTag('in.valid')

    with self.assertRaisesRegex(ValueError, 'Invalid tag'):
      ee.data.setDefaultWorkloadTag('Invalid')

    with self.assertRaisesRegex(ValueError, 'Invalid tag'):
      ee.data.setDefaultWorkloadTag('-invalid')

    with self.assertRaisesRegex(ValueError, 'Invalid tag'):
      ee.data.setDefaultWorkloadTag('invalid_')

    with self.assertRaisesRegex(ValueError, 'Invalid tag'):
      ee.data.setDefaultWorkloadTag('i' * 64)

    ee.data.setDefaultWorkloadTag('default-tag')
    self.assertEqual('default-tag', ee.data.getWorkloadTag())

    ee.data.setWorkloadTag('exports-1')
    self.assertEqual('exports-1', ee.data.getWorkloadTag())

    ee.data.setWorkloadTag('exports-2')
    self.assertEqual('exports-2', ee.data.getWorkloadTag())

    ee.data.resetWorkloadTag()
    self.assertEqual('default-tag', ee.data.getWorkloadTag())

    with ee.data.workloadTagContext('in-context'):
      self.assertEqual('in-context', ee.data.getWorkloadTag())

    self.assertEqual('default-tag', ee.data.getWorkloadTag())

    ee.data.setWorkloadTag('reset-me')
    self.assertEqual('reset-me', ee.data.getWorkloadTag())

    ee.data.setWorkloadTag('')
    self.assertEqual('', ee.data.getWorkloadTag())

    ee.data.setDefaultWorkloadTag('reset-me')
    self.assertEqual('reset-me', ee.data.getWorkloadTag())

    ee.data.resetWorkloadTag(True)
    self.assertEqual('', ee.data.getWorkloadTag())


def DoCloudProfileStubHttp(test, expect_profiling):

  def MockRequest(self, method, uri, data, headers, timeout):
    del self  # Unused.
    del method, uri, data, timeout  # Unused
    test.assertEqual(expect_profiling, ee.data._PROFILE_REQUEST_HEADER
                     in headers)
    response = requests.Response()
    response.status_code = 200
    response._content = '{"data": "dummy_data"}'
    if expect_profiling:
      response.headers[
          ee.data._PROFILE_RESPONSE_HEADER_LOWERCASE] = 'someProfileId'
    return response

  return mock.patch.object(requests.Session, 'request', new=MockRequest)


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