66import yaml
77import sys
88from os import path as osp
9+ from os import mkdir
910import gzip
1011import http
1112import json
13+ import tempfile
14+ import shutil
15+ import time
1216
1317def debug_http ():
1418 old_send = http .client .HTTPConnection .send
@@ -50,6 +54,18 @@ def load_data(path):
5054 else :
5155 return open (path , 'rb' ).read ()
5256
57+ def run_request_with_retry (run_req , expected_status_code = None , num_retries = 10 , wait_time = 0.5 ):
58+ for try_number in range (num_retries + 1 ):
59+ r = run_req ()
60+ if expected_status_code is None or r .status_code == expected_status_code :
61+ return r
62+ print ("Failed with" , r .text , r .status_code )
63+ if try_number < num_retries :
64+ print ("Retrying..." )
65+ time .sleep (wait_time )
66+ raise Exception ("Wrong status code. Got %s, expected %s" % (r .status_code , expected_status_code ))
67+
68+
5369def run_request_step (method , step ):
5470 assert method in {"GET" , "POST" , "PUT" , "DELETE" }
5571 if "headers" not in step :
@@ -70,12 +86,10 @@ def run_request_step(method, step):
7086 if ndjson is not None :
7187 kvargs ["data" ] = "\n " .join ([json .dumps (doc ) for doc in ndjson ])
7288 kvargs .setdefault ("headers" )["Content-Type" ] = "application/json"
73- r = method_req (url , ** kvargs )
74- expected_status_code = step .get ("status_code" , 200 )
75- if expected_status_code is not None :
76- if r .status_code != expected_status_code :
77- print (r .text )
78- raise Exception ("Wrong status code. Got %s, expected %s" % (r .status_code , expected_status_code ))
89+ expected_status_code = step .get ("status_code" , None )
90+ num_retries = step .get ("num_retries" , 0 )
91+ run_req = lambda : method_req (url , ** kvargs )
92+ r = run_request_with_retry (run_req , expected_status_code , num_retries )
7993 expected_resp = step .get ("expected" , None )
8094 if expected_resp is not None :
8195 try :
@@ -140,13 +154,15 @@ def add_path(self, path):
140154 path_tree .add_script (path_segs [- 1 ])
141155
142156 def visit_nodes (self , visitor , path = []):
143- visitor .enter_directory (path )
157+ success = True
158+ success &= visitor .enter_directory (path )
144159 for script in self .scripts :
145- visitor .run_scenario (path , script )
160+ success &= visitor .run_scenario (path , script )
146161 for k in sorted (self .children .keys ()):
147162 child_path = path + [k ]
148- self .children [k ].visit_nodes (visitor , child_path )
149- visitor .exit_directory (path )
163+ success &= self .children [k ].visit_nodes (visitor , child_path )
164+ success &= visitor .exit_directory (path )
165+ return success
150166
151167# Returns a new dictionary without modifying the arguments.
152168# The new dictionary is the result of merging the two dictionaries
@@ -164,10 +180,12 @@ def __init__(self, engine):
164180 self .context = {}
165181 def run_setup_teardown_scripts (self , script_name , path ):
166182 cwd = "/" .join (path )
183+ success = True
167184 for file_name in [script_name + ".yaml" , script_name + "." + self .engine + ".yaml" ]:
168185 script_fullpath = cwd + "/" + file_name
169186 if osp .exists (script_fullpath ):
170- self .run_scenario (path , file_name )
187+ success &= self .run_scenario (path , file_name )
188+ return success
171189 def load_context (self , path ):
172190 context = {"cwd" : "/" .join (path )}
173191 for file_name in ["_ctx.yaml" , "_ctx." + self .engine + ".yaml" ]:
@@ -180,13 +198,14 @@ def load_context(self, path):
180198 def enter_directory (self , path ):
181199 print ("============" )
182200 self .load_context (path )
183- self .run_setup_teardown_scripts ("_setup" , path )
201+ return self .run_setup_teardown_scripts ("_setup" , path )
184202 def exit_directory (self , path ):
185- self .run_setup_teardown_scripts ("_teardown" , path )
203+ success = self .run_setup_teardown_scripts ("_teardown" , path )
186204 self .context_stack .pop ()
187205 self .context = {}
188206 for ctx in self .context_stack :
189207 self .context .update (ctx )
208+ return success
190209 def run_scenario (self , path , script ):
191210 scenario_path = "/" .join (path + [script ])
192211 steps = list (open_scenario (scenario_path ))
@@ -208,9 +227,10 @@ def run_scenario(self, path, script):
208227 print (step )
209228 print (e )
210229 print ("--------------" )
211- break
230+ return False
212231 else :
213232 print ("🟢 %s: %d steps (%d skipped)" % (scenario_path , num_steps_executed , num_steps_skipped ))
233+ return True
214234
215235def build_path_tree (paths ):
216236 paths .sort ()
@@ -222,7 +242,7 @@ def build_path_tree(paths):
222242def run (scenario_paths , engine ):
223243 path_tree = build_path_tree (scenario_paths )
224244 visitor = Visitor (engine = engine )
225- path_tree .visit_nodes (visitor )
245+ return path_tree .visit_nodes (visitor )
226246
227247def filter_test (prefixes , test_name ):
228248 for prefix in prefixes :
@@ -239,6 +259,36 @@ def filter_tests(prefixes, test_names):
239259 if filter_test (prefixes , test_name )
240260 ]
241261
262+ class QuickwitRunner :
263+ def __init__ (self , quickwit_bin_path ):
264+ self .quickwit_dir = tempfile .TemporaryDirectory ()
265+ print ('created temporary directory' , self .quickwit_dir , self .quickwit_dir .name )
266+ qwdata = osp .join (self .quickwit_dir .name , "qwdata" )
267+ config = osp .join (self .quickwit_dir .name , "config" )
268+ mkdir (qwdata )
269+ mkdir (config )
270+ shutil .copy ("../../config/quickwit.yaml" , config )
271+ shutil .copy (quickwit_bin_path , self .quickwit_dir .name )
272+ self .proc = subprocess .Popen (["./quickwit" , "run" ], stdout = subprocess .DEVNULL , stderr = subprocess .DEVNULL , cwd = self .quickwit_dir .name )
273+ for i in range (100 ):
274+ try :
275+ print ("Checking on quickwit" )
276+ res = requests .get ("http://localhost:7280/health/readyz" )
277+ if res .status_code == 200 and res .text .strip () == "true" :
278+ print ("Quickwit started" )
279+ time .sleep (6 )
280+ break
281+ except :
282+ pass
283+ print ("Server not ready yet. Sleep and retry..." )
284+ time .sleep (1 )
285+ else :
286+ print ("Quickwit never started. Exiting." )
287+ sys .exit (2 )
288+ def __del__ (self ):
289+ print ("Killing Quickwit" )
290+ subprocess .Popen .kill (self .proc )
291+
242292def main ():
243293 import argparse
244294 arg_parser = argparse .ArgumentParser (
@@ -247,11 +297,28 @@ def main():
247297 )
248298 arg_parser .add_argument ("--engine" , help = "Targetted engine (elastic/quickwit)." , default = "quickwit" )
249299 arg_parser .add_argument ("--test" , help = "Specific prefix to select the tests to run. If not specified, all tests are run." , nargs = "*" )
300+ arg_parser .add_argument ("--binary" , help = "Specific the quickwit binary to run." , nargs = "?" )
250301 parsed_args = arg_parser .parse_args ()
302+
303+ print (parsed_args )
304+
305+ quickwit_process = None
306+ if parsed_args .binary is not None :
307+ if parsed_args .engine != "quickwit" :
308+ print ("The --binary option is only supported for quickwit engine." )
309+ sys .exit (3 )
310+ binary = parsed_args .binary
311+ quickwit_process = QuickwitRunner (binary )
312+ quickwit_process
313+
251314 scenario_filepaths = glob .glob ("scenarii/**/*.yaml" , recursive = True )
252315 scenario_filepaths = list (filter_tests (parsed_args .test , scenario_filepaths ))
253- run (scenario_filepaths , engine = parsed_args .engine )
316+ return run (scenario_filepaths , engine = parsed_args .engine )
254317
255318if __name__ == "__main__" :
256- main ()
319+ import sys
320+ if main ():
321+ sys .exit (0 )
322+ else :
323+ sys .exit (1 )
257324
0 commit comments