ClickHouse/dbms/tests/clickhouse-test

#!/usr/bin/env python
import sys
import os
import os.path
import re
import lxml.etree as et

from argparse import ArgumentParser
from argparse import FileType
from pprint import pprint
from subprocess import check_call
from subprocess import Popen
from subprocess import PIPE
from subprocess import CalledProcessError

if sys.stdout.isatty():
	COLORS = {
		"COLOR_RESET": "\033[0m",
		"COLOR_WHITE": "\033[1;37m",
		"COLOR_FAIL": "\033[1;31m",
		"COLOR_UNKNOWN": "\033[1;30m",
		"COLOR_OK": "\033[1;32m",
		"COLOR_SKIPPED": "\033[1;34m" }
else:
	COLORS = {
		"COLOR_RESET": "",
		"COLOR_WHITE": "",
		"COLOR_FAIL": "",
		"COLOR_UNKNOWN": "",
		"COLOR_OK": "",
		"COLOR_SKIPPED": "" }

MSG_FAIL = "{COLOR_WHITE}[ {COLOR_FAIL}FAIL{COLOR_WHITE} ]{COLOR_RESET}".format(**COLORS)
MSG_UNKNOWN = "{COLOR_WHITE}[ {COLOR_UNKNOWN}UNKNOWN{COLOR_WHITE} ]{COLOR_RESET}".format(**COLORS)
MSG_OK = "{COLOR_WHITE}[ {COLOR_OK}OK{COLOR_WHITE} ]{COLOR_RESET}".format(**COLORS)
MSG_SKIPPED = "{COLOR_WHITE}[ {COLOR_SKIPPED}SKIPPED{COLOR_WHITE} ]{COLOR_RESET}".format(**COLORS)

def main(args):

	SERVER_DIED = False
	
	def is_data_present():
		proc = Popen(args.client, stdin=PIPE, stdout=PIPE, stderr=PIPE)
		(stdout, stderr) = proc.communicate("EXISTS TABLE test.hits")
		if proc.returncode != 0:
			raise CalledProcessError(proc.returncode, args.client, stderr)
		
		return stdout.startswith('1')
	
	if args.zookeeper is None:
		try:
			check_call(['grep', '-q', '<zookeeper', '/etc/clickhouse-server/config-preprocessed.xml'], )
			args.zookeeper = True
		except CalledProcessError:
			args.zookeeper = False
	
	base_dir = os.path.abspath(args.queries)
	report = et.Element("testsuites", attrib = {'name': 'ClickHouse Tests'})
	
	tests_total = 0
	failures_total = 0
	disabled_total = 0
		
	for suite in sorted(os.listdir(base_dir)):
		if SERVER_DIED:
			break
		
		suite_dir = os.path.join(base_dir, suite)
		suite = re.search('^[0-9]+_(.*)$', suite).group(1)
		if os.path.isdir(suite_dir):
			print("\nRunning {} tests.\n".format(suite))
			
			report_suite = et.Element("testsuite", attrib = {"name": suite})
			
			tests = 0
			failures = 0
			disabled = 0
			if 'stateful' in suite and not is_data_present():
				print("Won't run stateful tests because test data wasn't loaded. See README.txt.")
				continue

			for case in sorted(filter(lambda case: re.search(args.test, case) if args.test else True, os.listdir(suite_dir))):
				if SERVER_DIED:
					break
					
				case_file = os.path.join(suite_dir, case)
				if os.path.isfile(case_file) and (case.endswith('.sh') or case.endswith('.sql')):
					(name, ext) = os.path.splitext(case)
					report_testcase = et.Element("testcase", attrib = {"name": name})
					
					print "{0:70}".format(name + ": "),
					
					if not args.zookeeper and 'zookeeper' in name:
						report_testcase.append(et.Element("skipped", attrib = {"message": "no zookeeper"}))
						disabled = disabled + 1
						print(MSG_SKIPPED, " - no zookeeper")
					else:
						if ext == '.sql':
							proc = Popen([args.client, '--multiquery'], stdin=PIPE, stdout = PIPE, stderr = PIPE)
							with open(case_file, 'r') as query_file:
								query = query_file.read()
								(stdout, stderr) = proc.communicate(query)
						else:
							proc = Popen(case_file, shell = True, stdout = PIPE, stderr = PIPE)
							(stdout, stderr) = proc.communicate()
						
						reference_file = os.path.join(suite_dir, name) + '.reference'
						
						if proc.returncode != 0:
							failure = et.Element("failure", attrib = {"message": "return code {}".format(proc.returncode)})
							report_testcase.append(failure)
							
							stdout_element = et.Element("system-out")
							stdout_element.text = et.CDATA(stdout)
							report_testcase.append(stdout_element)
							
							failures = failures + 1
							print("{0} - return code {1}".format(MSG_FAIL, proc.returncode))
							
							if stderr:
								stderr_element = et.Element("system-err")
								stderr_element.text = et.CDATA(stderr)
								report_testcase.append(stderr_element)
								print(stderr)
						
							if 'Connection refused' in stderr or 'Attempt to read after eof' in stderr:
								SERVER_DIED = True
								
						elif stderr:
							failure = et.Element("failure", attrib = {"message": "having stderror"})
							report_testcase.append(failure)
							
							stderr_element = et.Element("system-err")
							stderr_element.text = et.CDATA(stderr)
							report_testcase.append(stderr_element)
							
							failures = failures + 1
							print("{0} - having stderror:\n{1}".format(MSG_FAIL, stderr))
						elif 'Exception' in stdout:
							failure = et.Element("error", attrib = {"message": "having exception"})
							report_testcase.append(failure)
							
							stdout_element = et.Element("system-out")
							stdout_element.text = et.CDATA(stdout)
							report_testcase.append(stdout_element)
							
							failures = failures + 1
							print("{0} - having exception:\n{1}".format(MSG_FAIL, stdout))
						elif not os.path.isfile(reference_file):
							skipped = et.Element("skipped", attrib = {"message": "no reference file"})
							report_testcase.append(skipped)
							disabled = disabled + 1
							print("{0} - no reference file".format(MSG_UNKNOWN))
						else:
							result_file = os.path.join(suite_dir, name) + '.result'
							with open(result_file, 'w') as result:
								result.write(stdout)
							
							(diff, _) = Popen(['diff', reference_file, result_file], stdout = PIPE).communicate()
							
							if diff:
								failure = et.Element("failure", attrib = {"message": "result differs with reference"})
								report_testcase.append(failure)
								
								stdout_element = et.Element("system-out")
								stdout_element.text = et.CDATA(diff)
								report_testcase.append(stdout_element)
						
								failures = failures + 1
								print("{0} - result differs with reference:\n{1}".format(MSG_FAIL, diff))
							else:
								print(MSG_OK)
								os.remove(result_file)
					
					tests = tests + 1
					report_suite.append(report_testcase)
					
			report_suite.set("tests", str(tests))
			report_suite.set("failures", str(failures))
			report_suite.set("disabled", str(disabled))
			report_suite.set("skipped", str(disabled))
			report.append(report_suite)
			
			tests_total = tests_total + tests
			failures_total = failures_total + failures
			disabled_total = disabled_total + disabled
	
	report.set("tests", str(tests_total))
	report.set("failures", str(failures_total))
	report.set("disabled", str(disabled_total))
	
	if args.output is not None:
		args.output.write(et.tostring(report, encoding = "UTF-8", xml_declaration=True, pretty_print=True))

	if failures_total > 0:
		print("\n{COLOR_FAIL}Having {0} errors!{COLOR_RESET}".format(failures_total, **COLORS))
		sys.exit(1)
	else:
		print("\n{COLOR_OK}All tests passed.{COLOR_RESET}".format(**COLORS))
		sys.exit(0)
			

if __name__ == '__main__':
	parser = ArgumentParser(description = 'ClickHouse functional tests')
	parser.add_argument('-q', '--queries', default = 'queries', help = 'Path to queries dir')
	parser.add_argument('-c', '--client', default = 'clickhouse-client', help = 'Client program')
	parser.add_argument('-o', '--output', type = FileType('w'), help = 'Output xUnit compliant test report file')
	parser.add_argument('test', nargs = '?', help = 'Optional test case name regex')
	
	group = parser.add_mutually_exclusive_group(required = False)
	group.add_argument('--zookeeper', action = 'store_true', default = None, dest = 'zookeeper', help = 'Run zookeeper related tests')
	group.add_argument('--no-zookeeper', action = 'store_false', default = None, dest = 'zookeeper', help = 'Do not run zookeeper related tests')
	
	args = parser.parse_args()
	
	main(args)