diff --git a/src/Storages/StorageMergeTree.cpp b/src/Storages/StorageMergeTree.cpp
old mode 100644
new mode 100755
index 02c756d2e57..6db107afdc3
--- a/src/Storages/StorageMergeTree.cpp
+++ b/src/Storages/StorageMergeTree.cpp
@@ -781,17 +781,8 @@ std::shared_ptr<MergeMutateSelectedEntry> StorageMergeTree::selectOnePartitionTo
         if (part->modification_time < base)
             partition_parts_sum_diff[part->info.partition_id] += (base - part->modification_time);
     }
-    String best_partition_id;
-    Int32 max_diff = 0;
-    for (auto [partition_id, diff] : partition_parts_sum_diff)
-    {
-        if (diff > max_diff)
-        {
-           best_partition_id = partition_id;
-           max_diff = diff;
-        }
-    }
-    if (best_partition_id.empty())
+    auto best_partition_it = std::max_element(partition_parts_sum_diff.begin(), partition_parts_sum_diff.end(), [](const auto & e1, const auto & e2) { return e1.second < e2.second; });
+    if (best_partition_it == partition_parts_sum_diff.end())
     {
         return nullptr;
     }
@@ -801,7 +792,7 @@ std::shared_ptr<MergeMutateSelectedEntry> StorageMergeTree::selectOnePartitionTo
     auto merge_entry = selectPartsToMerge(
         metadata_snapshot,
         true,
-        best_partition_id,
+        best_partition_it->first,
         true,
         &disable_reason,
         table_lock_holder,
diff --git a/tests/integration/test_auto_optimize_partitions/test.py b/tests/integration/test_auto_optimize_partitions/test.py
index 138d1136696..6b0ec076fbe 100644
--- a/tests/integration/test_auto_optimize_partitions/test.py
+++ b/tests/integration/test_auto_optimize_partitions/test.py
@@ -21,6 +21,20 @@ def start_cluster():
     finally:
         cluster.shutdown()
 
+def check_expected_result_or_fail(seconds, expected):
+    ok = False
+    for i in range(int(seconds) * 2):
+        result = TSV(
+            node.query(
+                "SELECT count(*) FROM system.parts where table='test' and active=1"
+            )
+        )
+        if result == expected:
+            ok = True
+            break
+        else:
+            time.sleep(0.5)
+    assert(ok)
 
 def test_without_auto_optimize_merge_tree(start_cluster):
     node.query("CREATE TABLE test (i Int64) ENGINE = MergeTree ORDER BY i;")
@@ -28,17 +42,9 @@ def test_without_auto_optimize_merge_tree(start_cluster):
     node.query("INSERT INTO test SELECT 2")
     node.query("INSERT INTO test SELECT 3")
 
-    time.sleep(5)
 
     expected = TSV("""3\n""")
-    assert (
-        TSV(
-            node.query(
-                "SELECT count(*) FROM system.parts where table='test' and active=1"
-            )
-        )
-        == expected
-    )
+    check_expected_result_or_fail(5, expected)
 
     node.query("DROP TABLE test;")
 
@@ -51,17 +57,8 @@ def test_auto_optimize_merge_tree(start_cluster):
     node.query("INSERT INTO test SELECT 2")
     node.query("INSERT INTO test SELECT 3")
 
-    time.sleep(10)
-
     expected = TSV("""1\n""")
-    assert (
-        TSV(
-            node.query(
-                "SELECT count(*) FROM system.parts where table='test' and active=1"
-            )
-        )
-        == expected
-    )
+    check_expected_result_or_fail(10, expected)
 
     node.query("DROP TABLE test;")
 
@@ -74,16 +71,7 @@ def test_auto_optimize_replicated_merge_tree(start_cluster):
     node.query("INSERT INTO test SELECT 2")
     node.query("INSERT INTO test SELECT 3")
 
-    time.sleep(10)
-
     expected = TSV("""1\n""")
-    assert (
-        TSV(
-            node.query(
-                "SELECT count(*) FROM system.parts where table='test' and active=1"
-            )
-        )
-        == expected
-    )
+    check_expected_result_or_fail(10, expected)
 
     node.query("DROP TABLE test;")