From 05fc4fe8c42cfa12012b70c9573dc66e94b1befa Mon Sep 17 00:00:00 2001 From: Pierrick Rambaud Date: Sun, 30 Jun 2024 13:24:09 +0200 Subject: [PATCH] fix: use vanilla wait_for_task instead of a custom one --- pytest_gee/__init__.py | 7 ++++--- pytest_gee/utils.py | 24 +++++------------------- 2 files changed, 9 insertions(+), 22 deletions(-) diff --git a/pytest_gee/__init__.py b/pytest_gee/__init__.py index b765043..c86fa63 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,4 @@ 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) + return wait_for_task(task.id, timeout, log_progress=False) diff --git a/pytest_gee/utils.py b/pytest_gee/utils.py index 0485906..af3b4ee 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,7 @@ 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 + return wait_for_task(task.id, timeout, log_progress=False) def get_task(task_descripsion: str) -> Optional[ee.batch.Task]: @@ -118,7 +104,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)