From 389d783c4bab954bbc9fa548a16e12be7482c55e Mon Sep 17 00:00:00 2001
From: Rambaud Pierrick <12rambau@users.noreply.github.com>
Date: Tue, 13 Feb 2024 14:30:39 +0000
Subject: [PATCH] fix: solve the test issue

---
 pytest_gee/__init__.py | 34 ++++++++++++++++++++++++++++++++++
 pytest_gee/plugin.py   |  3 ++-
 tests/conftest.py      | 19 ++++++++++++++++++-
 3 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/pytest_gee/__init__.py b/pytest_gee/__init__.py
index 148f17e..06184b9 100644
--- a/pytest_gee/__init__.py
+++ b/pytest_gee/__init__.py
@@ -1,7 +1,9 @@
 """The init file of the package."""
 from __future__ import annotations
 
+import json
 import os
+import tempfile
 from pathlib import Path
 from typing import Union
 
@@ -42,6 +44,38 @@ def init_ee_from_token():
     ee.Initialize(http_transport=httplib2.Http())
 
 
+def init_ee_from_service_account():
+    """Initialize earth engine according using a service account.
+
+    The environment used to run the tests need to have a EARTHENGINE_SERVICE_ACCOUNT variable.
+    The content of this variable must be the copy of a personal credential file that you can generate from the google cloud console.
+
+    Note:
+        As all init method of ``pytest-gee``, this method will fallback to a regular ``ee.Initialize`` using the ``EARTHENGINE_PROJECT`` environment variable.
+    """
+    if "EARTHENGINE_SERVICE_ACCOUNT" in os.environ:
+        # extract the environment variables data
+        private_key = os.environ["EARTHENGINE_SERVICE_ACCOUNT"]
+        ee_user = json.loads(private_key)["client_email"]
+
+        # connect to GEE using a temp file to avoid writing the key to disk
+        with tempfile.TemporaryDirectory() as temp_dir:
+            file = Path(temp_dir) / "private_key.json"
+            file.write_text(private_key)
+            credentials = ee.ServiceAccountCredentials(ee_user, str(file))
+            ee.Initialize(credentials=credentials, http_transport=httplib2.Http())
+
+    elif "EARTHENGINE_PROJECT" in os.environ:
+        # if the user is in local development the authentication should already be available
+        # we simply need to use the provided project name
+        ee.Initialize(project=os.environ["EARTHENGINE_PROJECT"], http_transport=httplib2.Http())
+
+    else:
+        raise ValueError(
+            "EARTHENGINE_SERVICE_ACCOUNT or EARTHENGINE_PROJECT environment variable is missing"
+        )
+
+
 def wait(task: Union[ee.batch.Task, str], timeout: int = 5 * 60) -> str:
     """Wait until the selected process is finished or we reached timeout value.
 
diff --git a/pytest_gee/plugin.py b/pytest_gee/plugin.py
index 9a60864..b600a4d 100644
--- a/pytest_gee/plugin.py
+++ b/pytest_gee/plugin.py
@@ -19,7 +19,8 @@ def gee_hash():
 @pytest.fixture(scope="session")
 def gee_folder_root():
     """Link to the root folder of the connected account."""
-    return Path(ee.data.getAssetRoots()[0]["id"])
+    project_id = ee.data._cloud_api_user_project
+    return Path(f"projects/{project_id}/assets")
 
 
 @pytest.fixture(scope="session")
diff --git a/tests/conftest.py b/tests/conftest.py
index f72300e..85d03e0 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,4 +1,7 @@
 """Pytest session configuration."""
+import os
+from pathlib import Path
+
 import ee
 import pytest
 
@@ -7,7 +10,21 @@
 
 def pytest_configure():
     """Init GEE in the test environment."""
-    pytest_gee.init_ee_from_token()
+    pytest_gee.init_ee_from_service_account()
+
+
+# this fixture need override at the moment as the credential information cannot be reached from
+# the ee API as reported in https://issuetracker.google.com/issues/325020447
+@pytest.fixture(scope="session")
+def gee_folder_root():
+    """Link to the root folder of the connected account."""
+    project_id = os.environ.get("EARTHENGINE_PROJECT") or ee.data._cloud_api_user_project
+    if project_id is None:
+        raise ValueError(
+            "The project name cannot be detected."
+            "Please set the EARTHENGINE_PROJECT environment variable."
+        )
+    return Path(f"projects/{project_id}/assets")
 
 
 @pytest.fixture(scope="session")