fix: refactor run for NomosCli

This commit is contained in:
Roman 2025-02-13 10:29:18 +08:00
parent ebd56ce282
commit d340256605
No known key found for this signature in database
GPG Key ID: B8FE070B54E11B75
2 changed files with 37 additions and 29 deletions

View File

@ -1,6 +1,7 @@
import json
import os
import re
from enum import Enum
from src.data_storage import DS
from src.libs.common import generate_log_prefix
@ -61,33 +62,18 @@ class NomosCli:
DS.nomos_nodes.append(self)
@retry(stop=stop_after_delay(5), wait=wait_fixed(0.1), reraise=True)
def stop(self):
if self._container:
logger.debug(f"Stopping container with id {self._container.short_id}")
self._container.stop()
try:
self._container.remove()
except:
pass
self._container = None
logger.debug("Container stopped.")
match self._command:
case "reconstruct":
if "decode_only" in kwargs:
decode_only = kwargs["decode_only"]
else:
decode_only = False
return self.reconstruct(input_values=input_values, decode_only=decode_only)
case _:
return
@retry(stop=stop_after_delay(5), wait=wait_fixed(0.1), reraise=True)
def kill(self):
if self._container:
logger.debug(f"Killing container with id {self._container.short_id}")
self._container.kill()
try:
self._container.remove()
except:
pass
self._container = None
logger.debug("Container killed.")
def run_reconstruct(self, input_values=None, decode_only=False):
def reconstruct(self, input_values=None, decode_only=False):
keywords = ["Reconstructed data"]
self.run(input_values)
log_stream = self._container.logs(stream=True)
@ -112,3 +98,27 @@ class NomosCli:
DS.nomos_nodes.remove(self)
return result
@retry(stop=stop_after_delay(5), wait=wait_fixed(0.1), reraise=True)
def stop(self):
if self._container:
logger.debug(f"Stopping container with id {self._container.short_id}")
self._container.stop()
try:
self._container.remove()
except:
pass
self._container = None
logger.debug("Container stopped.")
@retry(stop=stop_after_delay(5), wait=wait_fixed(0.1), reraise=True)
def kill(self):
if self._container:
logger.debug(f"Killing container with id {self._container.short_id}")
self._container.kill()
try:
self._container.remove()
except:
pass
self._container = None
logger.debug("Container killed.")

View File

@ -25,8 +25,7 @@ class TestDataIntegrity(StepsDataAvailability):
rcv_data = self.get_data_range(selected_node, [0] * 31 + [1], [0] * 8, [0] * 7 + [5])
rcv_data_json = json.dumps(rcv_data)
cli = NomosCli(command="reconstruct")
reconstructed_data = cli.run_reconstruct(input_values=[rcv_data_json])
reconstructed_data = NomosCli(command="reconstruct").run(input_values=[rcv_data_json])
assert DATA_TO_DISPERSE[0] == reconstructed_data, "Reconstructed data are not same with original data"
@ -38,7 +37,6 @@ class TestDataIntegrity(StepsDataAvailability):
rcv_data = self.get_data_range(self.node2, [0] * 31 + [1], [0] * 8, [0] * 7 + [5])
rcv_data_json = json.dumps(rcv_data)
cli = NomosCli(command="reconstruct")
decoded_data = cli.run_reconstruct(input_values=[rcv_data_json], decode_only=True)
decoded_data = NomosCli(command="reconstruct").run(input_values=[rcv_data_json], decode_only=True)
assert DATA_TO_DISPERSE[0] == decoded_data, "Retrieved data are not same with original data"