ClickHouse/tests/integration/helpers/test_tools.py

119 lines
4.6 KiB
Python

import difflib
import time
import logging
from io import IOBase
class TSV:
"""Helper to get pretty diffs between expected and actual tab-separated value files"""
def __init__(self, contents):
if isinstance(contents, IOBase):
raw_lines = contents.readlines()
elif isinstance(contents, str) or isinstance(contents, str):
raw_lines = contents.splitlines(True)
elif isinstance(contents, list):
raw_lines = ['\t'.join(map(str, l)) if isinstance(l, list) else str(l) for l in contents]
elif isinstance(contents, TSV):
self.lines = contents.lines
return
else:
raise TypeError("contents must be either file or string or list, actual type: " + type(contents).__name__)
self.lines = [l.strip() for l in raw_lines if l.strip()]
def __eq__(self, other):
if not isinstance(other, TSV):
return self == TSV(other)
return self.lines == other.lines
def __ne__(self, other):
if not isinstance(other, TSV):
return self != TSV(other)
return self.lines != other.lines
def diff(self, other, n1='', n2=''):
if not isinstance(other, TSV):
return self.diff(TSV(other), n1=n1, n2=n2)
return list(line.rstrip() for line in difflib.unified_diff(self.lines, other.lines, fromfile=n1, tofile=n2))[2:]
def __str__(self):
return '\n'.join(self.lines)
def __repr__(self):
return self.__str__()
def __len__(self):
return len(self.lines)
@staticmethod
def toMat(contents):
return [line.split("\t") for line in contents.split("\n") if line.strip()]
def assert_eq_with_retry(instance, query, expectation, retry_count=20, sleep_time=0.5, stdin=None, timeout=None,
settings=None, user=None, ignore_error=False, get_result=lambda x: x):
expectation_tsv = TSV(expectation)
for i in range(retry_count):
try:
if TSV(get_result(instance.query(query, user=user, stdin=stdin, timeout=timeout, settings=settings,
ignore_error=ignore_error))) == expectation_tsv:
break
time.sleep(sleep_time)
except Exception as ex:
logging.exception(f"assert_eq_with_retry retry {i+1} exception {ex}")
time.sleep(sleep_time)
else:
val = TSV(get_result(instance.query(query, user=user, stdin=stdin, timeout=timeout, settings=settings,
ignore_error=ignore_error)))
if expectation_tsv != val:
raise AssertionError("'{}' != '{}'\n{}".format(expectation_tsv, val, '\n'.join(
expectation_tsv.diff(val, n1="expectation", n2="query"))))
def assert_logs_contain(instance, substring):
if not instance.contains_in_log(substring):
raise AssertionError("'{}' not found in logs".format(substring))
def assert_logs_contain_with_retry(instance, substring, retry_count=20, sleep_time=0.5):
for i in range(retry_count):
try:
if instance.contains_in_log(substring):
break
time.sleep(sleep_time)
except Exception as ex:
logging.exception(f"contains_in_log_with_retry retry {i+1} exception {ex}")
time.sleep(sleep_time)
else:
raise AssertionError("'{}' not found in logs".format(substring))
def exec_query_with_retry(instance, query, retry_count=40, sleep_time=0.5, silent=False, settings={}):
exception = None
for cnt in range(retry_count):
try:
res = instance.query(query, timeout=30, settings=settings)
if not silent:
logging.debug(f"Result of {query} on {cnt} try is {res}")
break
except Exception as ex:
exception = ex
if not silent:
logging.exception(f"Failed to execute query '{query}' on {cnt} try on instance '{instance.name}' will retry")
time.sleep(sleep_time)
else:
raise exception
def csv_compare(result, expected):
csv_result = TSV(result)
csv_expected = TSV(expected)
mismatch = []
max_len = len(csv_result) if len(csv_result) > len(csv_expected) else len(csv_expected)
for i in range(max_len):
if i >= len(csv_result):
mismatch.append("-[%d]=%s" % (i, csv_expected.lines[i]))
elif i >= len(csv_expected):
mismatch.append("+[%d]=%s" % (i, csv_result.lines[i]))
elif csv_expected.lines[i] != csv_result.lines[i]:
mismatch.append("-[%d]=%s" % (i, csv_expected.lines[i]))
mismatch.append("+[%d]=%s" % (i, csv_result.lines[i]))
return "\n".join(mismatch)