diff --git a/rapyuta_io/clients/device.py b/rapyuta_io/clients/device.py index 1333db7..e6f2fcb 100644 --- a/rapyuta_io/clients/device.py +++ b/rapyuta_io/clients/device.py @@ -573,9 +573,9 @@ def execute_command(self, command, retry_limit=0): jid = execution_result.get('jid') if not jid: raise ValueError("Job ID not found in the response") - return self.fetch_command_result(jid, [self.uuid], timeout=300) + return self.fetch_command_result(jid, [self.uuid], timeout=command.timeout) - def fetch_command_result(self, jid: str, deviceids: list, timeout: int = 300): + def fetch_command_result(self, jid: str, deviceids: list, timeout: int): """ Fetch the result of the command execution using the job ID (jid) and the first device ID from the list. Args: diff --git a/rapyuta_io/clients/model.py b/rapyuta_io/clients/model.py index 14e6f28..1af941f 100644 --- a/rapyuta_io/clients/model.py +++ b/rapyuta_io/clients/model.py @@ -61,7 +61,7 @@ class Command(ObjDict): """ - def __init__(self, cmd, shell=None, env=None, bg=False, runas=None, pwd=None, cwd=None): + def __init__(self, cmd, shell=None, env=None, bg=False, runas=None, pwd=None, cwd=None, timeout=300): super(ObjDict, self).__init__() if env is None: env = dict() @@ -73,6 +73,7 @@ def __init__(self, cmd, shell=None, env=None, bg=False, runas=None, pwd=None, cw self.cwd = pwd if cwd is not None: self.cwd = cwd + self.timeout = timeout self.validate() def validate(self): @@ -93,6 +94,8 @@ def validate(self): raise InvalidCommandException('Invalid environment variables') return raise InvalidCommandException('Invalid environment variables') + if self.timeout <= 0: + raise InvalidCommandException("Invalid timeout value") def to_json(self): # TODO: we need to rewrite this function.