address comments

facebookincubator · Dec 30, 2024 · c691bc3 · c691bc3
1 parent 6e475e2
commit c691bc3
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 9 deletions.
diff --git a/velox/connectors/hive/storage_adapters/abfs/AbfsConfig.cpp b/velox/connectors/hive/storage_adapters/abfs/AbfsConfig.cpp
@@ -23,7 +23,8 @@
 
 namespace facebook::velox::filesystems {
 
-std::unique_ptr<AzureDataLakeFileClient> AbfsConfig::fakeWriteClient_;
+std::function<std::unique_ptr<AzureDataLakeFileClient>()>
+    AbfsConfig::testWriteClientFn_;
 
 class DataLakeFileClientWrapper final : public AzureDataLakeFileClient {
  public:
@@ -161,8 +162,8 @@ std::unique_ptr<BlobClient> AbfsConfig::getReadFileClient() {
 }
 
 std::unique_ptr<AzureDataLakeFileClient> AbfsConfig::getWriteFileClient() {
-  if (fakeWriteClient_) {
-    return std::move(fakeWriteClient_);
+  if (testWriteClientFn_) {
+    return testWriteClientFn_();
   }
   std::unique_ptr<DataLakeFileClient> client;
   if (authType_ == kAzureSASAuthType) {

diff --git a/velox/connectors/hive/storage_adapters/abfs/AbfsConfig.h b/velox/connectors/hive/storage_adapters/abfs/AbfsConfig.h
@@ -94,9 +94,14 @@ class AbfsConfig {
   }
 
   /// Test only.
-  static void registerFakeWriteFileClient(
-      std::unique_ptr<AzureDataLakeFileClient> fakeClient) {
-    fakeWriteClient_ = std::move(fakeClient);
+  static void setUpTestWriteClient(
+      std::function<std::unique_ptr<AzureDataLakeFileClient>()> testClientFn) {
+    testWriteClientFn_ = testClientFn;
+  }
+
+  /// Test only.
+  static void tearDownTestWriteClient() {
+    testWriteClientFn_ = nullptr;
   }
 
  private:
@@ -118,7 +123,8 @@ class AbfsConfig {
   std::string authorityHost_;
   std::shared_ptr<Azure::Core::Credentials::TokenCredential> tokenCredential_;
 
-  static std::unique_ptr<AzureDataLakeFileClient> fakeWriteClient_;
+  static std::function<std::unique_ptr<AzureDataLakeFileClient>()>
+      testWriteClientFn_;
 };
 
 } // namespace facebook::velox::filesystems
diff --git a/velox/connectors/hive/storage_adapters/abfs/tests/AbfsFileSystemTest.cpp b/velox/connectors/hive/storage_adapters/abfs/tests/AbfsFileSystemTest.cpp
@@ -50,6 +50,8 @@ class AbfsFileSystemTest : public testing::Test {
   }
 
   void SetUp() override {
+    AbfsConfig::setUpTestWriteClient(
+        []() { return std::make_unique<MockDataLakeFileClient>(); });
     auto port = facebook::velox::exec::test::getFreePort();
     azuriteServer_ = std::make_shared<AzuriteServer>(port);
     azuriteServer_->start();
@@ -59,6 +61,7 @@ class AbfsFileSystemTest : public testing::Test {
   }
 
   void TearDown() override {
+    AbfsConfig::tearDownTestWriteClient();
     azuriteServer_->stop();
   }
 
@@ -277,8 +280,6 @@ TEST_F(AbfsFileSystemTest, registerAbfsFileSink) {
   auto hiveConfig =
       std::make_shared<const config::ConfigBase>(std::move(config));
   for (const auto& path : paths) {
-    AbfsConfig::registerFakeWriteFileClient(
-        std::make_unique<MockDataLakeFileClient>());
     auto sink = dwio::common::FileSink::create(
         path, {.connectorProperties = hiveConfig});
     auto writeFileSink = dynamic_cast<dwio::common::WriteFileSink*>(sink.get());