data: Update ee.data.create_assets to work with project-owned assets.

PiperOrigin-RevId: 740407354

Author: Nate Schmitz

python/ee/data.py

@@ -2331,16 +2331,17 @@ def create_assets(
     asset_ids: Sequence[str], asset_type: str, mk_parents: bool
 ) -> None:
   """Creates the specified assets if they do not exist."""
-  for asset_id in asset_ids:
+  ids = [convert_asset_id_to_asset_name(asset_id) for asset_id in asset_ids]
+  for asset_id in ids:
     if getInfo(asset_id):
       print('Asset %s already exists.' % asset_id)
       continue
     if mk_parents:
       parts = asset_id.split('/')
       # We don't need to create the namespace and the user's/project's folder.
-      if len(parts) > 2:
-        path = parts[0] + '/' + parts[1] + '/'
-        for part in parts[2:-1]:
+      if len(parts) > 3:
+        path = '/'.join(parts[0:3]) + '/'
+        for part in parts[3:-1]:
           path += part
           if getInfo(path) is None:
             createAsset({'type': ASSET_TYPE_FOLDER_CLOUD}, path)

python/ee/tests/data_test.py

@@ -1,8 +1,11 @@
 #!/usr/bin/env python3
 """Test for the ee.data module."""
 
+import json
+from typing import Any, Optional
 from unittest import mock
 
+import googleapiclient
 import httplib2
 import requests
 
@@ -14,6 +17,25 @@
 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: Optional[dict[str, int]] = None
+) -> dict[str, Any]:
+  return {
+      'type': 'FOLDER',
+      'name': name,
+      'quota': quota or {},
+  }
+
+
 class DataTest(unittest.TestCase):
 
   def setUp(self):
@@ -183,6 +205,66 @@ def testCreateFolder(self):
       asset = mock_create_asset.call_args.kwargs['body']
       self.assertEqual(asset, {'type': 'FOLDER'})
 
+  def testCreateAssets(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 testCreateAssets_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 testCreateAssets_noOpIfAssetExists(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 testCreateAssets_withParents(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 testStartIngestion(self):
     cloud_api_resource = mock.MagicMock()
     with apitestcase.UsingCloudApi(cloud_api_resource=cloud_api_resource):