YQ WM fixed flaky tests (ydb-platform#10805)

qyryq · Oct 24, 2024 · 188b6a2 · 188b6a2
1 parent 697d264
commit 188b6a2
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 7 deletions.
diff --git a/.github/config/muted_ya.txt b/.github/config/muted_ya.txt
@@ -44,7 +44,6 @@ ydb/core/kqp/ut/scheme [*/*]+chunk+chunk
 ydb/core/kqp/ut/service KqpService.CloseSessionsWithLoad
 ydb/core/kqp/ut/service [*/*] chunk chunk
 ydb/core/kqp/ut/service [*/*]+chunk+chunk
-ydb/core/kqp/workload_service/ut ResourcePoolClassifiersDdl.TestCreateResourcePoolClassifierOnServerless
 ydb/core/mind/hive/ut THiveTest.DrainWithHiveRestart
 ydb/core/persqueue/ut [*/*] chunk chunk
 ydb/core/quoter/ut QuoterWithKesusTest.PrefetchCoefficient

diff --git a/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp b/ydb/core/kqp/ut/scheme/kqp_scheme_ut.cpp
@@ -7365,7 +7365,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
             .Create();
 
         const auto& serverlessTenant = ydb->GetSettings().GetServerlessTenantName();
-        NWorkload::TSampleQueries::CheckSuccess(ydb->ExecuteQuery(R"(
+        ydb->ExecuteQueryRetry("Wait EnableResourcePoolsOnServerless", R"(
             CREATE RESOURCE POOL CLASSIFIER MyResourcePoolClassifier WITH (
                 RANK=20,
                 RESOURCE_POOL="test_pool"
@@ -7374,7 +7374,7 @@ Y_UNIT_TEST_SUITE(KqpScheme) {
                 .PoolId("")
                 .Database(serverlessTenant)
                 .NodeIndex(1)
-        ));
+        );
 
         const auto pathId = ydb->FetchDatabase(serverlessTenant)->Get()->PathId;
         UNIT_ASSERT_VALUES_EQUAL(

diff --git a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.cpp
@@ -432,6 +432,14 @@ class TWorkloadServiceYdbSetup : public IYdbSetup {
         return {.AsyncResult = promise.GetFuture(), .QueryRunnerActor = runerActor, .EdgeActor = edgeActor};
     }
 
+    void ExecuteQueryRetry(const TString& retryMessage, const TString& query, TQueryRunnerSettings settings = TQueryRunnerSettings(), TDuration timeout = FUTURE_WAIT_TIMEOUT) const override {
+        WaitFor(timeout, retryMessage, [this, query, settings](TString& errorString) {
+            auto result = ExecuteQuery(query, settings);
+            errorString = result.GetIssues().ToOneLineString();
+            return result.GetStatus() == EStatus::SUCCESS;
+        });
+    }
+
     // Async query execution actions
     void WaitQueryExecution(const TQueryRunnerResultAsync& query, TDuration timeout = FUTURE_WAIT_TIMEOUT) const override {
         auto event = GetRuntime()->GrabEdgeEvent<TEvQueryRunner::TEvExecutionStarted>(query.EdgeActor, timeout);

diff --git a/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h b/ydb/core/kqp/workload_service/ut/common/kqp_workload_service_ut_common.h
@@ -104,6 +104,7 @@ class IYdbSetup : public TThrRefBase {
     // Generic query helpers
     virtual TQueryRunnerResult ExecuteQuery(const TString& query, TQueryRunnerSettings settings = TQueryRunnerSettings()) const = 0;
     virtual TQueryRunnerResultAsync ExecuteQueryAsync(const TString& query, TQueryRunnerSettings settings = TQueryRunnerSettings()) const = 0;
+    virtual void ExecuteQueryRetry(const TString& retryMessage, const TString& query, TQueryRunnerSettings settings = TQueryRunnerSettings(), TDuration timeout = FUTURE_WAIT_TIMEOUT) const = 0;
 
     // Async query execution actions
     virtual void WaitQueryExecution(const TQueryRunnerResultAsync& query, TDuration timeout = FUTURE_WAIT_TIMEOUT) const = 0;

diff --git a/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp b/ydb/core/kqp/workload_service/ut/kqp_workload_service_ut.cpp
@@ -435,12 +435,12 @@ Y_UNIT_TEST_SUITE(ResourcePoolsDdl) {
             .NodeIndex(1);
 
         const TString& poolId = "my_pool";
-        TSampleQueries::CheckSuccess(ydb->ExecuteQuery(TStringBuilder() << R"(
+        ydb->ExecuteQueryRetry("Wait EnableResourcePoolsOnServerless", TStringBuilder() << R"(
             CREATE RESOURCE POOL )" << poolId << R"( WITH (
                 CONCURRENT_QUERY_LIMIT=1,
                 QUEUE_SIZE=0
             );
-        )", settings));
+        )", settings);
         settings.PoolId(poolId);
 
         auto hangingRequest = ydb->ExecuteQueryAsync(TSampleQueries::TSelect42::Query, settings.HangUpDuringExecution(true));
@@ -666,7 +666,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) {
     }
 
     void CreateSampleResourcePoolClassifier(TIntrusivePtr<IYdbSetup> ydb, const TString& classifierId, const TQueryRunnerSettings& settings, const TString& poolId) {
-        TSampleQueries::CheckSuccess(ydb->ExecuteQuery(TStringBuilder() << R"(
+        ydb->ExecuteQueryRetry("Wait EnableResourcePoolsOnServerless", TStringBuilder() << R"(
             GRANT ALL ON `)" << CanonizePath(settings.Database_ ? settings.Database_ : ydb->GetSettings().DomainName_) << R"(` TO `)" << settings.UserSID_ << R"(`;
             CREATE RESOURCE POOL )" << poolId << R"( WITH (
                 CONCURRENT_QUERY_LIMIT=0
@@ -680,7 +680,7 @@ Y_UNIT_TEST_SUITE(ResourcePoolClassifiersDdl) {
             .Database(settings.Database_)
             .NodeIndex(settings.NodeIndex_)
             .PoolId(NResourcePool::DEFAULT_POOL_ID)
-        ));
+        );
     }
 
     TString CreateSampleResourcePoolClassifier(TIntrusivePtr<IYdbSetup> ydb, const TQueryRunnerSettings& settings, const TString& poolId) {