Skip to content

Commit

Permalink
fix: solve the test issue
Browse files Browse the repository at this point in the history
  • Loading branch information
12rambau committed Feb 13, 2024
1 parent d251437 commit 389d783
Show file tree
Hide file tree
Showing 3 changed files with 54 additions and 2 deletions.
34 changes: 34 additions & 0 deletions pytest_gee/__init__.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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.
Expand Down
3 changes: 2 additions & 1 deletion pytest_gee/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
19 changes: 18 additions & 1 deletion tests/conftest.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
"""Pytest session configuration."""
import os
from pathlib import Path

import ee
import pytest

Expand All @@ -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")
Expand Down

0 comments on commit 389d783

Please sign in to comment.