Speedup testing

Signed-off-by: JaySon-Huang <jayson.hjs@gmail.com>

dbms/src/Flash/Disaggregated/tests/gtest_s3_lock_service.cpp

@@ -15,10 +15,10 @@
 #include <Common/Logger.h>
 #include <Common/typeid_cast.h>
 #include <Flash/Disaggregated/S3LockService.h>
+#include <Storages/S3/MockS3Client.h>
 #include <Storages/S3/S3Common.h>
 #include <Storages/S3/S3Filename.h>
 #include <Storages/tests/TiFlashStorageTestBasic.h>
-#include <TestUtils/MockS3Client.h>
 #include <TestUtils/TiFlashTestEnv.h>
 #include <TiDB/MockOwnerManager.h>
 #include <TiDB/OwnerManager.h>

dbms/src/TestUtils/MockS3Client.cpp → dbms/src/Storages/S3/MockS3Client.cpp

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 #include <Common/StringUtils/StringUtils.h>
-#include <TestUtils/MockS3Client.h>
+#include <Storages/S3/MockS3Client.h>
 #include <aws/core/AmazonWebServiceRequest.h>
 #include <aws/core/AmazonWebServiceResult.h>
 #include <aws/core/utils/xml/XmlSerializer.h>
@@ -27,24 +27,25 @@
 #include <aws/s3/model/Object.h>
 #include <aws/s3/model/PutObjectRequest.h>
 
-namespace DB
+namespace DB::S3::tests
 {
+using namespace Aws::S3;
 
-Aws::S3::Model::PutObjectOutcome MockS3Client::PutObject(const Aws::S3::Model::PutObjectRequest & r) const
+Model::PutObjectOutcome MockS3Client::PutObject(const Model::PutObjectRequest & r) const
 {
     put_keys.emplace_back(r.GetKey());
-    return Aws::S3::Model::PutObjectOutcome{Aws::AmazonWebServiceResult<Aws::Utils::Xml::XmlDocument>{}};
+    return Model::PutObjectOutcome{Aws::AmazonWebServiceResult<Aws::Utils::Xml::XmlDocument>{}};
 }
 
-Aws::S3::Model::DeleteObjectOutcome MockS3Client::DeleteObject(const Aws::S3::Model::DeleteObjectRequest & r) const
+Model::DeleteObjectOutcome MockS3Client::DeleteObject(const Model::DeleteObjectRequest & r) const
 {
     delete_keys.emplace_back(r.GetKey());
-    return Aws::S3::Model::DeleteObjectOutcome{Aws::AmazonWebServiceResult<Aws::Utils::Xml::XmlDocument>{}};
+    return Model::DeleteObjectOutcome{Aws::AmazonWebServiceResult<Aws::Utils::Xml::XmlDocument>{}};
 }
 
-Aws::S3::Model::ListObjectsV2Outcome MockS3Client::ListObjectsV2(const Aws::S3::Model::ListObjectsV2Request & r) const
+Model::ListObjectsV2Outcome MockS3Client::ListObjectsV2(const Model::ListObjectsV2Request & r) const
 {
-    Aws::S3::Model::ListObjectsV2Result resp;
+    Model::ListObjectsV2Result resp;
     for (const auto & k : put_keys)
     {
         if (startsWith(k, r.GetPrefix()))
@@ -60,7 +61,7 @@ Aws::S3::Model::ListObjectsV2Outcome MockS3Client::ListObjectsV2(const Aws::S3::
             }
             if (is_deleted)
                 continue;
-            Aws::S3::Model::Object o;
+            Model::Object o;
             o.SetKey(k);
             resp.AddContents(o);
         }
@@ -80,33 +81,33 @@ Aws::S3::Model::ListObjectsV2Outcome MockS3Client::ListObjectsV2(const Aws::S3::
             }
             if (is_deleted)
                 continue;
-            Aws::S3::Model::Object o;
+            Model::Object o;
             o.SetKey(k);
             resp.AddContents(o);
         }
     }
-    return Aws::S3::Model::ListObjectsV2Outcome{resp};
+    return Model::ListObjectsV2Outcome{resp};
 }
 
-Aws::S3::Model::HeadObjectOutcome MockS3Client::HeadObject(const Aws::S3::Model::HeadObjectRequest & r) const
+Model::HeadObjectOutcome MockS3Client::HeadObject(const Model::HeadObjectRequest & r) const
 {
     for (const auto & k : put_keys)
     {
         if (r.GetKey() == k)
         {
-            Aws::S3::Model::HeadObjectResult resp;
-            return Aws::S3::Model::HeadObjectOutcome{resp};
+            Model::HeadObjectResult resp;
+            return Model::HeadObjectOutcome{resp};
         }
     }
 
     if (!head_result_mtime)
     {
-        Aws::Client::AWSError error(Aws::S3::S3Errors::NO_SUCH_KEY, false);
-        return Aws::S3::Model::HeadObjectOutcome{error};
+        Aws::Client::AWSError error(S3Errors::NO_SUCH_KEY, false);
+        return Model::HeadObjectOutcome{error};
     }
-    Aws::S3::Model::HeadObjectResult resp;
+    Model::HeadObjectResult resp;
     resp.SetLastModified(head_result_mtime.value());
-    return Aws::S3::Model::HeadObjectOutcome{resp};
+    return Model::HeadObjectOutcome{resp};
 }
 
 void MockS3Client::clear()
@@ -117,4 +118,4 @@ void MockS3Client::clear()
     head_result_mtime.reset();
 }
 
-} // namespace DB
+} // namespace DB::S3::tests

dbms/src/TestUtils/MockS3Client.h → dbms/src/Storages/S3/MockS3Client.h

@@ -18,7 +18,7 @@
 #include <aws/s3/S3Client.h>
 #include <common/defines.h>
 
-namespace DB
+namespace DB::S3::tests
 {
 class MockS3Client final : public S3::TiFlashS3Client
 {
@@ -43,4 +43,4 @@ class MockS3Client final : public S3::TiFlashS3Client
     std::optional<Aws::Utils::DateTime> head_result_mtime;
     Aws::S3::Model::HeadObjectOutcome HeadObject(const Aws::S3::Model::HeadObjectRequest & request) const override;
 };
-} // namespace DB
+} // namespace DB::S3::tests

dbms/src/Storages/S3/S3Common.cpp

@@ -15,6 +15,7 @@
 #include <Common/Exception.h>
 #include <Common/ProfileEvents.h>
 #include <Common/Stopwatch.h>
+#include <Storages/S3/MockS3Client.h>
 #include <Storages/S3/S3Common.h>
 #include <aws/core/auth/AWSCredentials.h>
 #include <aws/core/http/Scheme.h>
@@ -136,16 +137,21 @@ bool ClientFactory::isEnabled() const
     return config.isS3Enabled();
 }
 
-void ClientFactory::init(const StorageS3Config & config_)
+void ClientFactory::init(const StorageS3Config & config_, bool mock_s3_)
 {
     config = config_;
     Aws::InitAPI(aws_options);
     Aws::Utils::Logging::InitializeAWSLogging(std::make_shared<AWSLogger>());
-    shared_client = create();
+    shared_client = mock_s3_ ? std::make_unique<tests::MockS3Client>() : create();
+    if (!mock_s3_)
     {
         auto raw_client = create();
         shared_tiflash_client = std::make_shared<TiFlashS3Client>(config.bucket, std::move(*raw_client.release()));
     }
+    else
+    {
+        shared_tiflash_client = std::make_shared<TiFlashS3Client>(config.bucket, tests::MockS3Client());
+    }
 }
 
 void ClientFactory::shutdown()

dbms/src/Storages/S3/S3Common.h

@@ -58,7 +58,7 @@ class ClientFactory
 
     bool isEnabled() const;
 
-    void init(const StorageS3Config & config_);
+    void init(const StorageS3Config & config_, bool mock_s3_ = false);
     void shutdown();
 
     const String & bucket() const;

dbms/src/TestUtils/gtests_dbms_main.cpp

@@ -84,6 +84,7 @@ int main(int argc, char ** argv)
         .access_key_id = access_key_id,
         .secret_access_key = secret_access_key,
     };
+    Poco::Environment::set("AWS_EC2_METADATA_DISABLED", "true"); // disable to speedup testing
     DB::S3::ClientFactory::instance().init(s3config);
 
 #ifdef FIU_ENABLE