From f350b57f4ba4a2d7863a39e7704baae8f3042ee2 Mon Sep 17 00:00:00 2001 From: Karol Kostrzewa Date: Tue, 2 Feb 2021 11:38:57 +0100 Subject: [PATCH 1/2] speedup celery tests --- tests/celery_tests.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/celery_tests.py b/tests/celery_tests.py index 689eea814bea9..63f8a0e7423bb 100644 --- a/tests/celery_tests.py +++ b/tests/celery_tests.py @@ -210,7 +210,6 @@ def test_run_sync_query_cta_config(setup_sqllab, ctas_method): ) assert query.select_sql == get_select_star(tmp_table_name, schema=CTAS_SCHEMA_NAME) - time.sleep(CELERY_SLEEP_TIME) results = run_sql(query.select_sql) assert QueryStatus.SUCCESS == results["status"], result @@ -231,9 +230,12 @@ def test_run_async_query_cta_config(setup_sqllab, ctas_method): QUERY, cta=True, ctas_method=ctas_method, async_=True, tmp_table=tmp_table_name, ) - time.sleep(CELERY_SLEEP_TIME) + for _ in range(CELERY_SLEEP_TIME * 2): + time.sleep(0.5) + query = get_query_by_id(result["query"]["serverId"]) + if QueryStatus.SUCCESS == query.status: + break - query = get_query_by_id(result["query"]["serverId"]) assert QueryStatus.SUCCESS == query.status assert get_select_star(tmp_table_name, schema=CTAS_SCHEMA_NAME) == query.select_sql assert ( @@ -252,9 +254,12 @@ def test_run_async_cta_query(setup_sqllab, ctas_method): QUERY, cta=True, ctas_method=ctas_method, async_=True, tmp_table=table_name ) - time.sleep(CELERY_SLEEP_TIME) + for _ in range(CELERY_SLEEP_TIME * 2): + time.sleep(0.5) + query = get_query_by_id(result["query"]["serverId"]) + if QueryStatus.SUCCESS == query.status: + break - query = get_query_by_id(result["query"]["serverId"]) assert QueryStatus.SUCCESS == query.status assert get_select_star(table_name) in query.select_sql @@ -274,9 +279,12 @@ def test_run_async_cta_query_with_lower_limit(setup_sqllab, ctas_method): result = run_sql( QUERY, cta=True, ctas_method=ctas_method, async_=True, tmp_table=tmp_table ) - time.sleep(CELERY_SLEEP_TIME) + for _ in range(CELERY_SLEEP_TIME * 2): + time.sleep(0.5) + query = get_query_by_id(result["query"]["serverId"]) + if QueryStatus.SUCCESS == query.status: + break - query = get_query_by_id(result["query"]["serverId"]) assert QueryStatus.SUCCESS == query.status assert get_select_star(tmp_table) == query.select_sql From 084fba9f1a24928db38c9f044bcb6e3398cccceb Mon Sep 17 00:00:00 2001 From: Karol Kostrzewa Date: Wed, 3 Feb 2021 13:00:31 +0100 Subject: [PATCH 2/2] refactor wait for success --- tests/celery_tests.py | 27 ++++++++++++--------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/tests/celery_tests.py b/tests/celery_tests.py index 63f8a0e7423bb..13c7ac3a8f083 100644 --- a/tests/celery_tests.py +++ b/tests/celery_tests.py @@ -230,11 +230,7 @@ def test_run_async_query_cta_config(setup_sqllab, ctas_method): QUERY, cta=True, ctas_method=ctas_method, async_=True, tmp_table=tmp_table_name, ) - for _ in range(CELERY_SLEEP_TIME * 2): - time.sleep(0.5) - query = get_query_by_id(result["query"]["serverId"]) - if QueryStatus.SUCCESS == query.status: - break + query = wait_for_success(result) assert QueryStatus.SUCCESS == query.status assert get_select_star(tmp_table_name, schema=CTAS_SCHEMA_NAME) == query.select_sql @@ -254,11 +250,7 @@ def test_run_async_cta_query(setup_sqllab, ctas_method): QUERY, cta=True, ctas_method=ctas_method, async_=True, tmp_table=table_name ) - for _ in range(CELERY_SLEEP_TIME * 2): - time.sleep(0.5) - query = get_query_by_id(result["query"]["serverId"]) - if QueryStatus.SUCCESS == query.status: - break + query = wait_for_success(result) assert QueryStatus.SUCCESS == query.status assert get_select_star(table_name) in query.select_sql @@ -279,11 +271,7 @@ def test_run_async_cta_query_with_lower_limit(setup_sqllab, ctas_method): result = run_sql( QUERY, cta=True, ctas_method=ctas_method, async_=True, tmp_table=tmp_table ) - for _ in range(CELERY_SLEEP_TIME * 2): - time.sleep(0.5) - query = get_query_by_id(result["query"]["serverId"]) - if QueryStatus.SUCCESS == query.status: - break + query = wait_for_success(result) assert QueryStatus.SUCCESS == query.status @@ -437,3 +425,12 @@ def my_task(): def delete_tmp_view_or_table(name: str, db_object_type: str): db.get_engine().execute(f"DROP {db_object_type} IF EXISTS {name}") + + +def wait_for_success(result): + for _ in range(CELERY_SLEEP_TIME * 2): + time.sleep(0.5) + query = get_query_by_id(result["query"]["serverId"]) + if QueryStatus.SUCCESS == query.status: + break + return query