def convert_df_to_string(
df: Union[pd.DataFrame, pd.Series],
n_rows: Optional[int] = None,
title: Optional[str] = None,
index: bool = False,
) -> str:
"""Convert DataFrame or Series to string for verifying test results.
:param df: DataFrame to be verified
:param n_rows: number of rows in expected output
:param title: title for test output
:return: string representation of input
"""
if isinstance(df, pd.Series):
df = df.to_frame()
n_rows = n_rows or len(df)
output = []
# Add title in the beginning if provided.
if title is not None:
output.append(prnt.frame(title))
# Provide context for full representation of data.
with pd.option_context(
"display.max_colwidth",
int(1e6),
"display.max_columns",
None,
"display.max_rows",
None,
):
# Add top N rows.
output.append(df.head(n_rows).to_string(index=index))
output_str = "\n".join(output)
return output_str
def diff_files(file_name1: str,
file_name2: str,
tag: Optional[str] = None) -> NoReturn:
# Diff to screen.
_, res = si.system_to_string(
"echo; sdiff -l -w 150 %s %s" % (file_name1, file_name2),
abort_on_error=False,
log_level=logging.DEBUG,
)
if tag is not None:
_LOG.error("%s", "\n" + prnt.frame(tag))
_LOG.error(res)
# Report how to diff.
vimdiff_cmd = "vimdiff %s %s" % (
os.path.abspath(file_name1),
os.path.abspath(file_name2),
)
# Save a script to diff.
diff_script = "./tmp_diff.sh"
io_.to_file(diff_script, vimdiff_cmd)
cmd = "chmod +x " + diff_script
si.system(cmd)
msg = []
msg.append("Diff with:")
msg.append("> " + vimdiff_cmd)
msg.append("or running:")
msg.append("> " + diff_script)
msg_as_str = "\n".join(msg)
_LOG.error(msg_as_str)
raise RuntimeError(msg_as_str)
def setUp(self) -> None:
random.seed(20000101)
np.random.seed(20000101)
# Name of the dir with artifacts for this test.
self._scratch_dir: Optional[str] = None
# Print banner to signal starting of a new test.
func_name = "%s.%s" % (self.__class__.__name__, self._testMethodName)
_LOG.debug("\n%s", prnt.frame(func_name))
def _assert_equal(
actual: str,
expected: str,
full_test_name: str,
test_dir: str,
fuzzy_match: bool = False,
) -> None:
"""Implement a better version of self.assertEqual() that reports
mismatching strings with sdiff and save them to files for further analysis
with vimdiff.
:param fuzzy: ignore differences in spaces and end of lines (see
`_remove_spaces`)
"""
def _to_string(obj: str) -> str:
if isinstance(obj, dict):
ret = pprint.pformat(obj)
else:
ret = str(obj)
ret = ret.rstrip("\n")
return ret
# Convert to strings.
actual = _to_string(actual)
expected = _to_string(expected)
# Fuzzy match, if needed.
if fuzzy_match:
_LOG.debug("Useing fuzzy match")
actual_orig = actual
actual = _remove_spaces(actual)
expected_orig = expected
expected = _remove_spaces(expected)
# Check.
if expected != actual:
_LOG.info(
"%s",
"\n" + prnt.frame("Test %s failed" % full_test_name, "=", 80))
if fuzzy_match:
# Set the following var to True to print the purified version (e.g.,
# tables too large).
print_purified_version = False
# print_purified_version = True
if print_purified_version:
expected = expected_orig
actual = actual_orig
# Dump the actual and expected strings to files.
_LOG.debug("Actual:\n%s", actual)
act_file_name = "%s/tmp.actual.txt" % test_dir
io_.to_file(act_file_name, actual)
_LOG.debug("Expected:\n%s", expected)
exp_file_name = "%s/tmp.expected.txt" % test_dir
io_.to_file(exp_file_name, expected)
#
tag = "ACTUAL vs EXPECTED"
diff_files(act_file_name, exp_file_name, tag)
def convert_info_to_string(info: Mapping) -> str:
"""Convert info to string for verifying test results.
Info often contains pd.Series, so pandas context is provided
to print all rows and all contents.
:param info: info to convert to string
:return: string representation of info
"""
output = []
# Provide context for full representation of pd.Series in info.
with pd.option_context(
"display.max_colwidth",
int(1e6),
"display.max_columns",
None,
"display.max_rows",
None,
):
output.append(prnt.frame("info"))
output.append(pprint.pformat(info))
output_str = "\n".join(output)
return output_str
def _system(
cmd: str,
abort_on_error: bool,
suppress_error: Optional[Any],
suppress_output: bool,
blocking: bool,
wrapper: Optional[Any],
output_file: Optional[Any],
tee: bool,
dry_run: bool,
log_level: Union[int, str],
) -> Tuple[int, str]:
"""Execute a shell command.
:param cmd: string with command to execute
:param abort_on_error: whether we should assert in case of error or not
:param suppress_error: set of error codes to suppress
:param suppress_output: whether to print the output or not
- If "on_debug_level" then print the output if the log level is DEBUG
:param blocking: blocking system call or not
:param wrapper: another command to prepend the execution of cmd
:param output_file: redirect stdout and stderr to this file
:param tee: if True, tee stdout and stderr to output_file
:param dry_run: just print the final command but not execute it
:param log_level: print the command to execute at level "log_level".
- If "echo" then print the command line to screen as print and not
logging
:return: return code (int), output of the command (str)
"""
orig_cmd = cmd[:]
# Prepare the command line.
cmd = "(%s)" % cmd
dbg.dassert_imply(tee, output_file is not None)
if output_file is not None:
dir_name = os.path.dirname(output_file)
if not os.path.exists(dir_name):
_LOG.debug("Dir '%s' doesn't exist: creating", dir_name)
dbg.dassert(bool(dir_name), "dir_name='%s'", dir_name)
os.makedirs(dir_name)
if tee:
cmd += " 2>&1 | tee %s" % output_file
else:
cmd += " 2>&1 >%s" % output_file
else:
cmd += " 2>&1"
if wrapper:
cmd = wrapper + " && " + cmd
#
# TODO(gp): Add a check for the valid values.
# TODO(gp): Make it "ECHO".
if isinstance(log_level, str):
dbg.dassert_eq(log_level, "echo")
print("> %s" % orig_cmd)
_LOG.debug("> %s", cmd)
else:
_LOG.log(log_level, "> %s", cmd)
#
dbg.dassert_in(suppress_output, ("ON_DEBUG_LEVEL", True, False))
if suppress_output == "ON_DEBUG_LEVEL":
# print("eff_lev=%s" % eff_level)
# print("lev=%s" % logging.DEBUG)
_LOG.getEffectiveLevel()
# Suppress the output if the verbosity level is higher than DEBUG,
# otherwise print.
suppress_output = _LOG.getEffectiveLevel() > logging.DEBUG
#
output = ""
if dry_run:
_LOG.warning("Not executing cmd\n%s\nas per user request", cmd)
rc = 0
return rc, output
# Execute the command.
try:
stdout = subprocess.PIPE
stderr = subprocess.STDOUT
p = subprocess.Popen(cmd,
shell=True,
executable="/bin/bash",
stdout=stdout,
stderr=stderr)
output = ""
if blocking:
# Blocking call: get the output.
while True:
line = p.stdout.readline().decode("utf-8") # type: ignore
if not line:
break
if not suppress_output:
print((line.rstrip("\n")))
output += line
p.stdout.close() # type: ignore
rc = p.wait()
else:
# Not blocking.
# Wait until process terminates (without using p.wait()).
max_cnt = 20
cnt = 0
while p.poll() is None:
# Process hasn't exited yet, let's wait some time.
time.sleep(0.1)
cnt += 1
_LOG.debug("cnt=%s, rc=%s", cnt, p.returncode)
if cnt > max_cnt:
break
if cnt > max_cnt:
# Timeout: we assume it worked.
rc = 0
else:
rc = p.returncode
if suppress_error is not None:
dbg.dassert_isinstance(suppress_error, set)
if rc in suppress_error:
rc = 0
except OSError as e:
rc = -1
_LOG.error("error=%s", str(e))
_LOG.debug("rc=%s", rc)
if abort_on_error and rc != 0:
msg = ("\n" + prnt.frame("cmd='%s' failed with rc='%s'" % (cmd, rc)) +
"\nOutput of the failing command is:\n%s\n%s\n%s" %
(prnt.line(">"), output, prnt.line("<")))
_LOG.error("%s", msg)
raise RuntimeError("cmd='%s' failed with rc='%s'" % (cmd, rc))
# dbg.dassert_type_in(output, (str, ))
return rc, output