diff --git a/pytest_gee/__init__.py b/pytest_gee/__init__.py index b765043..c037376 100644 --- a/pytest_gee/__init__.py +++ b/pytest_gee/__init__.py @@ -9,8 +9,8 @@ import ee import httplib2 - -from pytest_gee import utils +from deprecated.sphinx import deprecated +from ee.cli.utils import wait_for_task __version__ = "0.3.4" __author__ = "Pierrick Rambaud" @@ -83,6 +83,7 @@ def init_ee_from_service_account(): ) +@deprecated(version="0.3.5", reason="Use the vanilla GEE ``wait_for_task`` function instead.") 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. @@ -95,4 +96,5 @@ def wait(task: Union[ee.batch.Task, str], timeout: int = 5 * 60) -> str: """ # just expose the utils function # this is compulsory as wait is also needed in the utils module - return utils.wait(task, timeout) + task_id = task.id if isinstance(task, ee.batch.Task) else task + return wait_for_task(task_id, timeout, log_progress=False) diff --git a/pytest_gee/utils.py b/pytest_gee/utils.py index 0485906..2b64199 100644 --- a/pytest_gee/utils.py +++ b/pytest_gee/utils.py @@ -6,13 +6,15 @@ """ from __future__ import annotations -import time from pathlib import Path, PurePosixPath from typing import List, Optional, Union import ee +from deprecated.sphinx import deprecated +from ee.cli.utils import wait_for_task +@deprecated(version="0.3.5", reason="Use the vanilla GEE ``wait_for_task`` function instead.") def wait(task: Union[ee.batch.Task, str], timeout: int = 10 * 60) -> str: """Wait until the selected process is finished or we reached timeout value. @@ -23,23 +25,8 @@ def wait(task: Union[ee.batch.Task, str], timeout: int = 10 * 60) -> str: Returns: the final state of the task """ - # give 5 seconds of delay to GEE to make sure the task is created - time.sleep(5) - - # init both the task object and the state - task = task if isinstance(task, ee.batch.Task) else get_task(task) - assert task is not None, "The task is not found" - state = "UNSUBMITTED" - - # loop every 5s to check the task state. This is blocking the Python interpreter - start_time = time.time() - while state != "COMPLETED" and time.time() - start_time < timeout: - time.sleep(5) - state = task.status()["state"] - if state == "FAILED": - break - - return state + task_id = task.id if isinstance(task, ee.batch.Task) else task + return wait_for_task(task_id, timeout, log_progress=False) def get_task(task_descripsion: str) -> Optional[ee.batch.Task]: @@ -118,7 +105,7 @@ def export_asset( # launch the task and wait for the end of exportation task.start() - wait(description) + wait_for_task(task.id, 10 * 60, False) return PurePosixPath(asset_id)