Kurt Schwehr 4884055c0a data.py: Lint and add test coverage
PiperOrigin-RevId: 833400162
2025-11-17 10:34:46 -08:00

1358 lines
51 KiB
Python

#!/usr/bin/env python3
"""Test for the ee.data module."""
import json
from typing import Any
from unittest import mock
import googleapiclient
import httplib2
import requests
import unittest
import ee
from ee import _cloud_api_utils
from ee import _state
from ee import apitestcase
from ee import featurecollection
from ee import image
def NotFoundError() -> googleapiclient.errors.HttpError:
"""Creates a mock HttpError with a 404 status code."""
resp = httplib2.Response({'status': '404', 'reason': 'Not Found'})
content = json.dumps({'error': {'code': 404, 'message': 'Not Found'}}).encode(
'utf-8'
)
return googleapiclient.errors.HttpError(resp, content)
def NewFolderAsset(
name: str, quota: dict[str, int] | None = None
) -> dict[str, Any]:
return {
'type': 'FOLDER',
'name': name,
'quota': quota or {},
}
class DataTest(unittest.TestCase):
def setUp(self):
super().setUp()
mock.patch.object(
ee.data,
'getAlgorithms',
return_value=apitestcase.GetAlgorithms(),
autospec=True,
).start()
def tearDown(self):
super().tearDown()
ee.data.reset()
mock.patch.stopall()
def test_is_initialized(self):
self.assertFalse(ee.data.is_initialized())
with apitestcase.UsingCloudApi():
self.assertTrue(ee.data.is_initialized())
@mock.patch.object(ee.data, '_install_cloud_api_resource', return_value=None)
def test_initialize(self, mock_install_cloud_api_resource):
ee.data.initialize()
self.assertTrue(ee.data.is_initialized())
mock_install_cloud_api_resource.assert_called_once()
@mock.patch.object(ee.data, '_install_cloud_api_resource', return_value=None)
def test_initialize_with_project(
self, unused_mock_install_cloud_api_resource
):
ee.data.initialize(project='my-project')
self.assertTrue(ee.data.is_initialized())
self.assertEqual(
_state.get_state().cloud_api_user_project, 'my-project'
)
@mock.patch.object(ee.data, '_install_cloud_api_resource', return_value=None)
def test_initialize_with_no_project(
self, unused_mock_install_cloud_api_resource
):
ee.data.initialize()
self.assertTrue(ee.data.is_initialized())
self.assertEqual(
_state.get_state().cloud_api_user_project, 'earthengine-legacy'
)
@mock.patch.object(ee.data, '_install_cloud_api_resource', return_value=None)
def test_initialize_with_credentials(
self, unused_mock_install_cloud_api_resource
):
creds = mock.MagicMock()
ee.data.initialize(credentials=creds)
self.assertTrue(ee.data.is_initialized())
self.assertEqual(creds, _state.get_state().credentials)
@mock.patch.object(ee.data, '_install_cloud_api_resource', return_value=None)
def test_initialize_with_cloud_api_key(
self, unused_mock_install_cloud_api_resource
):
cloud_api_key = 'a cloud api key'
ee.data.initialize(cloud_api_key=cloud_api_key)
self.assertTrue(ee.data.is_initialized())
self.assertEqual(cloud_api_key, _state.get_state().cloud_api_key)
def test_set_max_retries_bad_values(self):
with self.assertRaises(ValueError):
ee.data.setMaxRetries(-1)
with self.assertRaises(ValueError):
ee.data.setMaxRetries(100)
def test_set_max_retries(self):
mock_result = {'result': 5}
ee.data.setMaxRetries(3)
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
cloud_api_resource.projects().value().compute().execute.return_value = (
mock_result
)
self.assertEqual(5, ee.data.computeValue(ee.Number(1)))
self.assertEqual(
3,
cloud_api_resource.projects()
.value()
.compute()
.execute.call_args.kwargs['num_retries'],
)
def test_set_cloud_api_key(self):
cloud_api_key = 'a cloud api key'
with mock.patch.object(
ee.data, '_install_cloud_api_resource', return_value=None
) as mock_install_cloud_api_resource:
ee.data.setCloudApiKey(cloud_api_key)
self.assertEqual(cloud_api_key, _state.get_state().cloud_api_key)
mock_install_cloud_api_resource.assert_called_once()
def test_set_deadline(self):
deadline_ms = 12345
with mock.patch.object(
ee.data, '_install_cloud_api_resource', return_value=None
) as mock_install_cloud_api_resource:
ee.data.setDeadline(deadline_ms)
self.assertEqual(deadline_ms, _state.get_state().deadline_ms)
mock_install_cloud_api_resource.assert_called_once()
def test_get_set_user_agent(self):
self.assertIsNone(ee.data.getUserAgent())
user_agent = 'user-agent'
ee.data.setUserAgent(user_agent)
self.assertEqual(user_agent, ee.data.getUserAgent())
def test_authorize_http_no_credentials(self):
self.assertIsNone(ee.data._get_state().credentials)
http = mock.MagicMock()
self.assertEqual(http, ee.data.authorizeHttp(http))
def test_authorize_http_with_credentials(self):
creds = mock.MagicMock()
ee.data._get_state().credentials = creds
http = mock.MagicMock()
with mock.patch.object(
ee.data.google_auth_httplib2, 'AuthorizedHttp'
) as mock_authorized_http:
result = ee.data.authorizeHttp(http)
self.assertEqual(mock_authorized_http.return_value, result)
mock_authorized_http.assert_called_once_with(creds)
def test_list_operations(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 test_list_operations_empty_list(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 test_get_operation(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
name = 'projects/test-project/operations/foo'
cloud_api_resource.projects().operations().get.execute.return_value = {
'name': name,
'done': False,
}
ee.data.getOperation(name)
cloud_api_resource.projects().operations().get.assert_called_once_with(
name=name
)
def test_get_task_status(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
cloud_api_resource.projects().operations().get.return_value.execute.return_value = {
'name': 'projects/earthengine-legacy/operations/foo',
'done': False,
'metadata': {'state': 'RUNNING'},
}
result = ee.data.getTaskStatus('foo')
cloud_api_resource.projects().operations().get.assert_called_once_with(
name='projects/earthengine-legacy/operations/foo'
)
self.assertEqual(
result,
[{
'id': 'foo',
'state': 'RUNNING',
'name': 'projects/earthengine-legacy/operations/foo',
}],
)
def test_get_task_status_with_not_found(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
cloud_api_resource.projects().operations().get.return_value.execute.side_effect = [
{
'name': 'projects/earthengine-legacy/operations/foo',
'done': False,
'metadata': {'state': 'RUNNING'},
},
NotFoundError(),
{
'name': 'projects/earthengine-legacy/operations/bar',
'done': True,
'metadata': {'state': 'SUCCEEDED'},
},
]
result = ee.data.getTaskStatus(['foo', 'missing', 'bar'])
cloud_api_resource.projects().operations().get.assert_has_calls([
mock.call(name='projects/earthengine-legacy/operations/foo'),
mock.call().execute(num_retries=5),
mock.call(name='projects/earthengine-legacy/operations/missing'),
mock.call().execute(num_retries=5),
mock.call(name='projects/earthengine-legacy/operations/bar'),
mock.call().execute(num_retries=5),
])
self.assertEqual(
3,
cloud_api_resource.projects()
.operations()
.get.return_value.execute.call_count,
)
self.assertEqual(
result,
[
{
'id': 'foo',
'state': 'RUNNING',
'name': 'projects/earthengine-legacy/operations/foo',
},
{
'id': 'missing',
'state': 'UNKNOWN',
'name': 'projects/earthengine-legacy/operations/missing',
},
{
'id': 'bar',
'state': 'COMPLETED',
'name': 'projects/earthengine-legacy/operations/bar',
},
],
)
def test_cancel_operation(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
cancel_mock = cloud_api_resource.projects().operations().cancel
cancel_mock.execute.return_value = {}
ee.data.cancelOperation('projects/test-project/operations/foo')
cancel_mock.assert_called_once_with(
name='projects/test-project/operations/foo', body={}
)
def test_cancel_task(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
cancel_mock = cloud_api_resource.projects().operations().cancel
cancel_mock.execute.return_value = {}
ee.data.cancelTask('foo')
cancel_mock.assert_called_once_with(
name='projects/earthengine-legacy/operations/foo', body={}
)
def test_create_asset(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 test_create_asset_with_v1alpha_params(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']},
)
def test_create_folder(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.createFolder('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'})
@mock.patch.object(ee.data, 'createAsset')
def test_create_asset_home(self, mock_create_asset):
ee.data.createAssetHome('users/test')
mock_create_asset.assert_called_once_with({
'name': 'projects/earthengine-legacy/assets/users/test',
'type': 'FOLDER',
})
def test_create_assets(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
asset_name = 'projects/some-project/assets/some-asset'
cloud_api_resource.projects().assets().get().execute.side_effect = (
NotFoundError()
)
ee.data.create_assets([asset_name], 'FOLDER', False)
mock_create_asset = cloud_api_resource.projects().assets().create
mock_create_asset.assert_called_once_with(
parent='projects/some-project',
assetId='some-asset',
body={'type': 'FOLDER'},
prettyPrint=False,
)
def test_create_assets_empty(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
ee.data.create_assets([], 'FOLDER', False)
mock_create_asset = cloud_api_resource.projects().assets().create
mock_create_asset.assert_not_called()
def test_create_assets_no_op_if_asset_exists(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
asset_name = 'projects/some-project/assets/some-asset'
cloud_api_resource.projects().assets().get.execute.return_value = (
NewFolderAsset(asset_name)
)
ee.data.create_assets([asset_name], 'FOLDER', False)
mock_create_asset = cloud_api_resource.projects().assets().create
mock_create_asset.assert_not_called()
def test_create_assets_with_parents(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
asset_name = 'projects/some-project/assets/foo/bar'
cloud_api_resource.projects().assets().get().execute.side_effect = (
NotFoundError()
)
ee.data.create_assets([asset_name], 'FOLDER', True)
mock_create_asset = cloud_api_resource.projects().assets().create
mock_create_asset.assert_has_calls([
mock.call(
parent='projects/some-project',
assetId='foo',
body={'type': 'FOLDER'},
prettyPrint=False,
),
mock.call().execute(num_retries=5),
mock.call(
parent='projects/some-project',
assetId='foo/bar',
body={'type': 'FOLDER'},
prettyPrint=False,
),
mock.call().execute(num_retries=5),
])
def test_start_ingestion(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
mock_result = {'name': 'operations/ingestion', 'done': False}
cloud_api_resource.projects().image().import_().execute.return_value = (
mock_result
)
manifest = {
'id': 'projects/some-project/assets/some-name',
'arg': 'something',
}
result = ee.data.startIngestion(
'request_id',
manifest,
True
)
self.assertEqual(result['id'], 'ingestion')
self.assertEqual(result['name'], 'operations/ingestion')
mock_import = cloud_api_resource.projects().image().import_
import_args = mock_import.call_args.kwargs['body']
self.assertEqual(
import_args['imageManifest'],
{
'name': 'projects/some-project/assets/some-name',
'arg': 'something',
},
)
self.assertTrue(import_args['overwrite'])
def test_start_table_ingestion(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
mock_result = {'name': 'operations/ingestion', 'done': False}
cloud_api_resource.projects().table().import_.return_value.execute.return_value = (
mock_result
)
params = {
'id': 'users/test/table',
'sources': [{'uris': ['gs://bucket/file.shp'], 'charset': 'UTF-8'}],
}
result = ee.data.startTableIngestion('request_id', params, True)
self.assertEqual(result['id'], 'ingestion')
self.assertEqual(result['name'], 'operations/ingestion')
mock_import = cloud_api_resource.projects().table().import_
mock_import.assert_called_once()
call_kwargs = mock_import.call_args.kwargs
self.assertEqual(call_kwargs['project'], 'projects/earthengine-legacy')
body = call_kwargs['body']
self.assertEqual(
body['tableManifest'],
{
'name': 'projects/earthengine-legacy/assets/users/test/table',
'sources': [
{'uris': ['gs://bucket/file.shp'], 'charset': 'UTF-8'}
],
},
)
self.assertEqual(body['requestId'], 'request_id')
self.assertTrue(body['overwrite'])
def test_start_external_image_ingestion(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
cloud_api_resource.projects().image().importExternal.return_value.execute.return_value = (
{}
)
manifest = {
'id': 'users/test/image',
'tilesets': [{'sources': [{'uris': ['gs://bucket/file.tif']}]}],
}
result = ee.data.startExternalImageIngestion(manifest, True)
expected_name = 'projects/earthengine-legacy/assets/users/test/image'
self.assertEqual(result['name'], expected_name)
mock_import_external = (
cloud_api_resource.projects().image().importExternal
)
mock_import_external.assert_called_once()
call_kwargs = mock_import_external.call_args.kwargs
self.assertEqual(call_kwargs['project'], 'projects/earthengine-legacy')
body = call_kwargs['body']
self.assertEqual(
body['imageManifest'],
{
'name': expected_name,
'tilesets': [{'sources': [{'uris': ['gs://bucket/file.tif']}]}],
},
)
self.assertTrue(body['overwrite'])
def test_set_asset_properties(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),
{'properties.\"mYPropErTy\"',
'properties.\"system:time_start\"'})
def test_update_asset(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
asset_id = 'users/test/asset'
asset = {'properties': {'foo': 'bar'}}
update_mask = ['properties.foo']
ee.data.updateAsset(asset_id, asset, update_mask)
cloud_api_resource.projects().assets().patch.assert_called_once_with(
name='projects/earthengine-legacy/assets/users/test/asset',
body={'updateMask': {'paths': update_mask}, 'asset': asset},
)
def test_list_assets(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('path/to/folder')
cloud_api_resource.projects().assets().listAssets(
).execute.assert_called_once()
self.assertEqual(mock_result, actual_result)
def test_list_assets_with_page_size(self):
mock_http = mock.MagicMock(httplib2.Http)
ok_resp = httplib2.Response({'status': 200})
page = (
b'{"assets": [{"path": "id1", "type": "type1"}], "nextPageToken": "t1"}'
)
mock_http.request.side_effect = [(ok_resp, page)]
with apitestcase.UsingCloudApi(mock_http=mock_http):
actual_result = ee.data.listAssets(
{'parent': 'path/to/folder', 'pageSize': 3}
)
expected_result = {
'assets': [{'path': 'id1', 'type': 'type1'}],
'nextPageToken': 't1',
}
self.assertEqual(expected_result, actual_result)
def test_list_assets_multiple_pages(self):
mock_http = mock.MagicMock(httplib2.Http)
ok_resp = httplib2.Response({'status': 200})
page1 = (
b'{"assets": [{"path": "id1", "type": "type1"}], "nextPageToken": "t1"}'
)
page2 = (
b'{"assets": [{"path": "id2", "type": "type2"}], "nextPageToken": "t2"}'
)
page3 = b'{"assets": [{"path": "id3", "type": "type3"}]}'
mock_http.request.side_effect = [
(ok_resp, page1),
(ok_resp, page2),
(ok_resp, page3),
]
with apitestcase.UsingCloudApi(mock_http=mock_http):
actual_result = ee.data.listAssets('path/to/folder')
expected_result = {
'assets': [
{'path': 'id1', 'type': 'type1'},
{'path': 'id2', 'type': 'type2'},
{'path': 'id3', 'type': 'type3'},
]
}
self.assertEqual(expected_result, actual_result)
def test_list_images(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('path/to/folder')
cloud_api_resource.projects().assets().listAssets(
).execute.assert_called_once()
self.assertEqual({'images': [{
'path': 'id1',
'type': 'type1'
}]}, actual_result)
def test_list_images_with_page_size(self):
mock_http = mock.MagicMock(httplib2.Http)
ok_resp = httplib2.Response({'status': 200})
page = (
b'{"assets": [{"path": "id1", "type": "type1"}], "nextPageToken": "t1"}'
)
mock_http.request.side_effect = [(ok_resp, page)]
with apitestcase.UsingCloudApi(mock_http=mock_http):
actual_result = ee.data.listImages(
{'parent': 'path/to/folder', 'pageSize': 3}
)
expected_result = {
'images': [{'path': 'id1', 'type': 'type1'}],
'nextPageToken': 't1',
}
self.assertEqual(expected_result, actual_result)
def test_list_images_multiple_pages(self):
mock_http = mock.MagicMock(httplib2.Http)
ok_resp = httplib2.Response({'status': 200})
page1 = (
b'{"assets": [{"path": "id1", "type": "type1"}], "nextPageToken": "t1"}'
)
page2 = (
b'{"assets": [{"path": "id2", "type": "type2"}], "nextPageToken": "t2"}'
)
page3 = b'{"assets": [{"path": "id3", "type": "type3"}]}'
mock_http.request.side_effect = [
(ok_resp, page1),
(ok_resp, page2),
(ok_resp, page3),
]
with apitestcase.UsingCloudApi(mock_http=mock_http):
actual_result = ee.data.listImages('path/to/folder')
expected_result = {
'images': [
{'path': 'id1', 'type': 'type1'},
{'path': 'id2', 'type': 'type2'},
{'path': 'id3', 'type': 'type3'},
]
}
self.assertEqual(expected_result, actual_result)
def test_list_buckets(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 test_get_asset_roots(self):
with mock.patch.object(
ee.data,
'listBuckets',
return_value={'assets': [{'name': 'id1', 'type': 'FOLDER'}]},
) as mock_list_buckets:
result = ee.data.getAssetRoots()
mock_list_buckets.assert_called_once()
self.assertEqual([{'id': 'id1', 'type': 'Folder'}], result)
def test_simple_get_list_via_cloud_api(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 test_get_list_asset_root_via_cloud_api(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 test_get_list_asset_root_via_cloud_api_no_slash(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 test_complex_get_list_via_cloud_api(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 test_get_map_id(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 test_get_map_id_with_workload_tag(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'])
def test_get_download_id(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 test_get_download_id_with_workload_tag(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 test_get_download_id_with_band_list(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 test_get_download_id_with_image_id(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 test_get_download_id_with_serialized_image(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 test_get_thumb_id(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 test_get_thumb_id_with_workload_tag(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 test_get_table_download_id(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 test_get_table_download_id_with_workload_tag(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 test_cloud_profiling_enabled(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 test_cloud_profiling_disabled(self):
with apitestcase.UsingCloudApi(), DoCloudProfileStubHttp(self, False):
ee.data.listImages({'parent': 'projects/earthengine-public/assets/q'})
def test_cloud_error_translation(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 test_list_features(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)
def test_get_pixels(self):
cloud_api_resource_raw = mock.MagicMock()
with apitestcase.UsingCloudApi(
cloud_api_resource_raw=cloud_api_resource_raw
):
assets = cloud_api_resource_raw.projects().assets()
mock_result = b'pixel data'
assets.getPixels.return_value.execute.return_value = mock_result
asset_id = 'users/foo/bar'
params = {'assetId': asset_id}
result = ee.data.getPixels(params)
self.assertEqual(mock_result, result)
assets.getPixels.assert_called_once_with(
name='projects/earthengine-legacy/assets/users/foo/bar',
body={'fileFormat': 'AUTO_JPEG_PNG'},
)
def test_compute_pixels(self):
cloud_api_resource_raw = mock.MagicMock()
with apitestcase.UsingCloudApi(
cloud_api_resource_raw=cloud_api_resource_raw
):
mock_result = b'pixel data'
(
cloud_api_resource_raw.projects()
.image()
.computePixels.return_value.execute.return_value
) = mock_result
expression = ee.Image(1)
params = {'expression': expression}
result = ee.data.computePixels(params)
self.assertEqual(mock_result, result)
(
cloud_api_resource_raw.projects()
.image()
.computePixels.assert_called_once_with(
project='projects/earthengine-legacy',
body={
'expression': ee.serializer.encode(expression),
'fileFormat': 'AUTO_JPEG_PNG',
},
)
)
def test_get_feature_view_tiles_key(self):
cloud_api_resource = mock.MagicMock()
_state.get_state().tile_base_url = 'base_url'
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 test_get_project_config(self) -> None:
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
mock_result = {'fake-project-config-value': 1}
cloud_api_resource.projects().getConfig().execute.return_value = (
mock_result
)
actual_result = ee.data.getProjectConfig()
cloud_api_resource.projects().getConfig().execute.assert_called_once()
self.assertEqual(mock_result, actual_result)
def test_update_project_config(self) -> None:
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
mock_result = {'fake-project-config-value': 1}
cloud_api_resource.projects().updateConfig().execute.return_value = (
mock_result
)
project_config = {'maxConcurrentExports': 2}
actual_result = ee.data.updateProjectConfig(
project_config, ['max_concurrent_exports']
)
cloud_api_resource.projects().updateConfig.assert_called_with(
name='projects/earthengine-legacy/config',
body=project_config,
updateMask='max_concurrent_exports',
)
cloud_api_resource.projects().updateConfig().execute.assert_called_once()
self.assertEqual(mock_result, actual_result)
def test_update_project_config_no_mask(self) -> None:
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
mock_result = {'fake-project-config-value': 1}
cloud_api_resource.projects().updateConfig().execute.return_value = (
mock_result
)
project_config = {'maxConcurrentExports': 2}
actual_result = ee.data.updateProjectConfig(project_config)
cloud_api_resource.projects().updateConfig.assert_called_with(
name='projects/earthengine-legacy/config',
body=project_config,
updateMask='max_concurrent_exports',
)
cloud_api_resource.projects().updateConfig().execute.assert_called_once()
self.assertEqual(mock_result, actual_result)
def test_workload_tag(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 test_reset_workload_tag_opt_params(self):
ee.data.setDefaultWorkloadTag('reset-me')
self.assertEqual('reset-me', ee.data.getWorkloadTag())
ee.data.resetWorkloadTag(opt_resetDefault=True)
self.assertEqual('', ee.data.getWorkloadTag())
def test_get_asset_root_quota_v1alpha(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
fake_asset = {
'type': 'FOLDER',
'name': 'projects/test-proj/assets',
'quota': {
'assetCount': 123,
'maxAssets': 456,
'sizeBytes': 789,
'maxSizeBytes': 1001,
},
}
cloud_api_resource.projects().assets().get().execute.return_value = (
fake_asset
)
quota = ee.data.getAssetRootQuota('projects/test-proj/assets')
expected = {
'asset_count': {'usage': 123, 'limit': 456},
'asset_size': {'usage': 789, 'limit': 1001},
}
self.assertEqual(expected, quota)
def test_get_asset_root_quota(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
fake_asset = {
'type': 'FOLDER',
'name': 'projects/test-proj/assets',
'quota': {
'assetCount': 123,
'maxAssetCount': 456,
'sizeBytes': 789,
'maxSizeBytes': 1001,
},
}
cloud_api_resource.projects().assets().get().execute.return_value = (
fake_asset
)
quota = ee.data.getAssetRootQuota('projects/test-proj/assets')
expected = {
'asset_count': {'usage': 123, 'limit': 456},
'asset_size': {'usage': 789, 'limit': 1001},
}
self.assertEqual(expected, quota)
def test_get_asset_root_quota_not_root(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
asset_id = 'users/test/asset'
fake_asset = {
'type': 'IMAGE',
'name': 'projects/earthengine-legacy/assets/users/test/asset',
}
cloud_api_resource.projects().assets().get().execute.return_value = (
fake_asset
)
with self.assertRaisesRegex(
ee.ee_exception.EEException, f'{asset_id} is not a root folder.'
):
ee.data.getAssetRootQuota(asset_id)
def test_get_iam_policy(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
asset_id = 'users/test/asset'
ee.data.getIamPolicy(asset_id)
cloud_api_resource.projects().assets().getIamPolicy.assert_called_once_with(
resource='projects/earthengine-legacy/assets/users/test/asset',
body={},
prettyPrint=False,
)
def test_get_asset_acl(self):
asset_id = 'users/test/asset'
policy = {'bindings': [{'role': 'roles/viewer', 'members': ['allUsers']}]}
acl = {'readers': ['allUsers']}
with mock.patch.object(
ee.data, 'getIamPolicy', return_value=policy
) as mock_get_iam_policy, mock.patch.object(
ee.data._cloud_api_utils, 'convert_iam_policy_to_acl', return_value=acl
) as mock_convert:
result = ee.data.getAssetAcl(asset_id)
mock_get_iam_policy.assert_called_once_with(asset_id)
mock_convert.assert_called_once_with(policy)
self.assertEqual(acl, result)
def test_set_asset_acl(self):
asset_id = 'users/test/asset'
acl_update_dict = {'readers': ['allUsers']}
acl_update_str = '{"readers": ["allUsers"]}'
policy = {'bindings': [{'role': 'roles/viewer', 'members': ['allUsers']}]}
with mock.patch.object(
ee.data._cloud_api_utils,
'convert_acl_to_iam_policy',
return_value=policy,
) as mock_convert, mock.patch.object(
ee.data, 'setIamPolicy'
) as mock_set_iam_policy:
ee.data.setAssetAcl(asset_id, acl_update_dict)
mock_convert.assert_called_once_with(acl_update_dict)
mock_set_iam_policy.assert_called_once_with(asset_id, policy)
mock_convert.reset_mock()
mock_set_iam_policy.reset_mock()
ee.data.setAssetAcl(asset_id, acl_update_str)
mock_convert.assert_called_once_with(acl_update_dict)
mock_set_iam_policy.assert_called_once_with(asset_id, policy)
def test_set_iam_policy(self):
cloud_api_resource = mock.MagicMock()
with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):
asset_id = 'users/test/asset'
policy = {'bindings': [{'role': 'roles/viewer', 'members': ['allUsers']}]}
ee.data.setIamPolicy(asset_id, policy)
cloud_api_resource.projects().assets().setIamPolicy.assert_called_once_with(
resource='projects/earthengine-legacy/assets/users/test/asset',
body={'policy': policy},
prettyPrint=False,
)
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()