diff --git a/tests/flow/tests_dag.py b/tests/flow/tests_dag.py index 01b030b47..a00c9d97e 100644 --- a/tests/flow/tests_dag.py +++ b/tests/flow/tests_dag.py @@ -761,6 +761,7 @@ def run(): 'AI.TENSORGET', result_tensor_keyname, 'META', ) + t.join() ensureSlaveSynced(con, env) env.assertEqual([b'OK',b'OK',[b'dtype', b'FLOAT', b'shape', [1, 2]]], ret) diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py index 39bd85b33..849920560 100644 --- a/tests/flow/tests_onnx.py +++ b/tests/flow/tests_onnx.py @@ -215,6 +215,7 @@ def run(): t.start() con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'a{1}', 'OUTPUTS', 'b{1}') + t.join() ensureSlaveSynced(con, env) diff --git a/tests/flow/tests_pytorch.py b/tests/flow/tests_pytorch.py index b38bc1d56..ed7c0f1b9 100644 --- a/tests/flow/tests_pytorch.py +++ b/tests/flow/tests_pytorch.py @@ -269,7 +269,8 @@ def test_pytorch_modelrun_autobatch_badbatch(env): def run(): con = env.getConnection() try: - con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'd{1}', 'e{1}', 'OUTPUTS', 'f1{1}', 'f2{1}') + ret = con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'd{1}', 'e{1}', 'OUTPUTS', 'f1{1}', 'f2{1}') + env.assertEqual(ret, b'OK') except Exception as e: exception = e env.assertEqual(type(exception), redis.exceptions.ResponseError) @@ -279,11 +280,13 @@ def run(): t.start() try: - con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'a{1}', 'b{1}', 'OUTPUTS', 'c1{1}', 'c2{1}') + ret = con.execute_command('AI.MODELRUN', 'm{1}', 'INPUTS', 'a{1}', 'b{1}', 'OUTPUTS', 'c1{1}', 'c2{1}') + env.assertEqual(ret, b'OK') except Exception as e: exception = e env.assertEqual(type(exception), redis.exceptions.ResponseError) env.assertEqual("Model did not generate the expected batch size", exception.__str__()) + t.join()