fix ut

Signed-off-by: srlch <[email protected]>
StarRocks · Aug 7, 2024 · 6ea02f0 · 6ea02f0
1 parent 9015834
commit 6ea02f0
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 4 deletions.
diff --git a/be/src/fs/fs_s3.cpp b/be/src/fs/fs_s3.cpp
@@ -112,7 +112,8 @@ class S3ClientFactory {
     }
 
     // Only use for UT
-    bool find_client_cache_keys_by_config_TEST(const Aws::Client::ClientConfiguration& config) {
+    bool find_client_cache_keys_by_config_TEST(const Aws::Client::ClientConfiguration& config,
+                                               AWSCloudConfiguration* cloud_config = nullptr) {
         return _find_client_cache_keys_by_config_TEST(config);
     }
 
@@ -135,9 +136,11 @@ class S3ClientFactory {
     constexpr static int kMaxItems = 8;
 
     // Only use for UT
-    bool _find_client_cache_keys_by_config_TEST(const Aws::Client::ClientConfiguration& config) {
+    bool _find_client_cache_keys_by_config_TEST(const Aws::Client::ClientConfiguration& config,
+                                                AWSCloudConfiguration* cloud_config = nullptr) {
         for (size_t i = 0; i < _items; i++) {
-            if (_client_cache_keys[i] == ClientCacheKey{config, AWSCloudConfiguration{}}) return true;
+            if (_client_cache_keys[i] == ClientCacheKey{config, cloud_config == nullptr ?
+                                                                AWSCloudConfiguration{} : *cloud_config}) return true;
         }
         return false;
     }
@@ -321,7 +324,7 @@ static std::shared_ptr<Aws::S3::S3Client> new_s3client(
         const TCloudConfiguration& tCloudConfiguration = (opts.cloud_configuration != nullptr)
                                                                  ? *opts.cloud_configuration
                                                                  : hdfs_properties->cloud_configuration;
-        return S3ClientFactory::instance().new_client(tCloudConfiguration);
+        return S3ClientFactory::instance().new_client(tCloudConfiguration, operation_type);
     } else if (hdfs_properties != nullptr) {
         DCHECK(hdfs_properties->__isset.end_point);
         if (hdfs_properties->__isset.end_point) {

diff --git a/be/test/fs/fs_s3_test.cpp b/be/test/fs/fs_s3_test.cpp
@@ -544,12 +544,37 @@ TEST_F(S3FileSystemTest, test_new_S3_client_with_rename_operation) {
     ASSERT_TRUE(S3ClientFactory::instance().find_client_cache_keys_by_config_TEST(config));
 
     // use config::object_storage_request_timeout_ms instead
+    int old_object_storage_rename_file_request_timeout_ms = config::object_storage_rename_file_request_timeout_ms;
+    int old_object_storage_request_timeout_ms = config::object_storage_request_timeout_ms;
     config::object_storage_rename_file_request_timeout_ms = -1;
     config::object_storage_request_timeout_ms = 1000;
     // only used for generate a new S3 client into global cache
     (void)fs->rename_file(S3Path("/dir/source_name"), S3Path("/dir/target_name"));
     config.requestTimeoutMs = config::object_storage_request_timeout_ms;
     ASSERT_TRUE(S3ClientFactory::instance().find_client_cache_keys_by_config_TEST(config));
+    config::object_storage_rename_file_request_timeout_ms = old_object_storage_rename_file_request_timeout_ms;
+    config::object_storage_request_timeout_ms = old_object_storage_request_timeout_ms;
+
+    std::map<std::string, std::string> test_properties;
+    test_properties[AWS_S3_USE_AWS_SDK_DEFAULT_BEHAVIOR] = "true";
+    TCloudConfiguration tCloudConfiguration;
+    tCloudConfiguration.__set_cloud_type(TCloudType::AWS);
+    tCloudConfiguration.__set_cloud_properties_v2(test_properties);
+    auto cloud_config = CloudConfigurationFactory::create_aws(tCloudConfiguration);
+
+    config.requestTimeoutMs = config::object_storage_rename_file_request_timeout_ms;
+    (void)S3ClientFactory::instance().new_client(tCloudConfiguration, S3ClientFactory::OperationType::RENAME_FILE);
+    ASSERT_TRUE(S3ClientFactory::instance().find_client_cache_keys_by_config_TEST(config, &cloud_config));
+
+    old_object_storage_rename_file_request_timeout_ms = config::object_storage_rename_file_request_timeout_ms;
+    old_object_storage_request_timeout_ms = config::object_storage_request_timeout_ms;
+    config::object_storage_rename_file_request_timeout_ms = -1;
+    config::object_storage_request_timeout_ms = 1000;
+    (void)S3ClientFactory::instance().new_client(tCloudConfiguration, S3ClientFactory::OperationType::RENAME_FILE);
+    config.requestTimeoutMs = config::object_storage_request_timeout_ms;
+    ASSERT_TRUE(S3ClientFactory::instance().find_client_cache_keys_by_config_TEST(config, &cloud_config));
+    config::object_storage_rename_file_request_timeout_ms = old_object_storage_rename_file_request_timeout_ms;
+    config::object_storage_request_timeout_ms = old_object_storage_request_timeout_ms;
 }
 
 } // namespace starrocks