fix(unittest): Init with dbfilename= before attempting to save (#…

…2127) This is a pretty recent regression.
dragonflydb · Nov 6, 2023 · efeae54 · efeae54
1 parent 7e23c14
commit efeae54
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 27 deletions.
diff --git a/src/server/cluster/cluster_family_test.cc b/src/server/cluster/cluster_family_test.cc
@@ -680,6 +680,8 @@ TEST_F(ClusterFamilyTest, ClusterModeSelectNotAllowed) {
 }
 
 TEST_F(ClusterFamilyTest, ClusterFirstConfigCallDropsEntriesNotOwnedByNode) {
+  InitWithDbFilename();
+
   Run({"debug", "populate", "50000"});
 
   EXPECT_EQ(Run({"save", "df"}), "OK");

diff --git a/src/server/generic_family_test.cc b/src/server/generic_family_test.cc
@@ -633,41 +633,44 @@ TEST_F(GenericFamilyTest, Restore) {
 }
 
 TEST_F(GenericFamilyTest, Info) {
+  InitWithDbFilename();  // Needed for `save`
+
   auto get_rdb_changes_since_last_save = [](const string& str) -> size_t {
     const string matcher = "rdb_changes_since_last_save:";
     const auto pos = str.find(matcher) + matcher.size();
     const auto sub = str.substr(pos, 1);
     return atoi(sub.c_str());
   };
-  auto resp = Run({"set", "k", "1"});
-  resp = Run({"info", "persistence"});
+
+  EXPECT_EQ(Run({"set", "k", "1"}), "OK");
+  auto resp = Run({"info", "persistence"});
   EXPECT_EQ(1, get_rdb_changes_since_last_save(resp.GetString()));
 
-  resp = Run({"set", "k", "1"});
+  EXPECT_EQ(Run({"set", "k", "1"}), "OK");
   resp = Run({"info", "persistence"});
   EXPECT_EQ(2, get_rdb_changes_since_last_save(resp.GetString()));
 
-  resp = Run({"set", "k2", "2"});
+  EXPECT_EQ(Run({"set", "k2", "2"}), "OK");
   resp = Run({"info", "persistence"});
   EXPECT_EQ(3, get_rdb_changes_since_last_save(resp.GetString()));
 
-  resp = Run({"save"});
+  EXPECT_EQ(Run({"save"}), "OK");
   resp = Run({"info", "persistence"});
   EXPECT_EQ(0, get_rdb_changes_since_last_save(resp.GetString()));
 
-  resp = Run({"set", "k2", "2"});
+  EXPECT_EQ(Run({"set", "k2", "2"}), "OK");
   resp = Run({"info", "persistence"});
   EXPECT_EQ(1, get_rdb_changes_since_last_save(resp.GetString()));
 
-  resp = Run({"bgsave"});
+  EXPECT_EQ(Run({"bgsave"}), "OK");
   resp = Run({"info", "persistence"});
   EXPECT_EQ(0, get_rdb_changes_since_last_save(resp.GetString()));
 
-  resp = Run({"set", "k3", "3"});
+  EXPECT_EQ(Run({"set", "k3", "3"}), "OK");
   resp = Run({"info", "persistence"});
   EXPECT_EQ(1, get_rdb_changes_since_last_save(resp.GetString()));
 
-  resp = Run({"del", "k3"});
+  EXPECT_THAT(Run({"del", "k3"}), IntArg(1));
   resp = Run({"info", "persistence"});
   EXPECT_EQ(1, get_rdb_changes_since_last_save(resp.GetString()));
 }

diff --git a/src/server/rdb_test.cc b/src/server/rdb_test.cc
@@ -18,7 +18,6 @@ extern "C" {
 #include "base/logging.h"
 #include "facade/facade_test.h"  // needed to find operator== for RespExpr.
 #include "io/file.h"
-#include "io/file_util.h"
 #include "server/engine_shard_set.h"
 #include "server/rdb_load.h"
 #include "server/test_utils.h"
@@ -33,33 +32,18 @@ using absl::StrCat;
 ABSL_DECLARE_FLAG(int32, list_compress_depth);
 ABSL_DECLARE_FLAG(int32, list_max_listpack_size);
 ABSL_DECLARE_FLAG(int, compression_mode);
-ABSL_DECLARE_FLAG(string, dbfilename);
 
 namespace dfly {
 
 class RdbTest : public BaseFamilyTest {
  protected:
-  void TearDown();
   void SetUp();
 
   io::FileSource GetSource(string name);
 };
 
 void RdbTest::SetUp() {
-  SetFlag(&FLAGS_dbfilename, "rdbtestdump");
-  BaseFamilyTest::SetUp();
-}
-
-void RdbTest::TearDown() {
-  // Disable save on shutdown
-  SetFlag(&FLAGS_dbfilename, "");
-
-  auto rdb_files = io::StatFiles("rdbtestdump*");
-  CHECK(rdb_files);
-  for (const auto& fl : *rdb_files) {
-    unlink(fl.name.c_str());
-  }
-  BaseFamilyTest::TearDown();
+  InitWithDbFilename();
 }
 
 inline const uint8_t* to_byte(const void* s) {

diff --git a/src/server/test_utils.cc b/src/server/test_utils.cc
@@ -19,6 +19,7 @@ extern "C" {
 #include "base/logging.h"
 #include "base/stl_util.h"
 #include "facade/dragonfly_connection.h"
+#include "io/file_util.h"
 #include "server/acl/acl_log.h"
 #include "util/fibers/pool.h"
 
@@ -257,7 +258,13 @@ void BaseFamilyTest::ResetService() {
 }
 
 void BaseFamilyTest::ShutdownService() {
-  DCHECK(service_);
+  if (service_ == nullptr) {
+    return;
+  }
+
+  // Don't save files during shutdown
+  CleanupSnapshots();
+  absl::SetFlag(&FLAGS_dbfilename, "");
 
   service_->Shutdown();
   service_.reset();
@@ -270,6 +277,26 @@ void BaseFamilyTest::ShutdownService() {
   pp_->Stop();
 }
 
+void BaseFamilyTest::InitWithDbFilename() {
+  ShutdownService();
+
+  absl::SetFlag(&FLAGS_dbfilename, "rdbtestdump");
+  CleanupSnapshots();
+  ResetService();
+}
+
+void BaseFamilyTest::CleanupSnapshots() {
+  string dbfilename = absl::GetFlag(FLAGS_dbfilename);
+  if (dbfilename.empty())
+    return;
+
+  auto rdb_files = io::StatFiles(absl::StrCat(dbfilename, "*"));
+  CHECK(rdb_files);
+  for (const auto& fl : *rdb_files) {
+    unlink(fl.name.c_str());
+  }
+}
+
 unsigned BaseFamilyTest::NumLocked() {
   atomic_uint count = 0;
   shard_set->RunBriefInParallel([&](EngineShard* shard) {

diff --git a/src/server/test_utils.h b/src/server/test_utils.h
@@ -94,6 +94,9 @@ class BaseFamilyTest : public ::testing::Test {
 
   void ShutdownService();
 
+  void InitWithDbFilename();
+  void CleanupSnapshots();
+
   bool IsLocked(DbIndex db_index, std::string_view key) const;
   ConnectionContext::DebugInfo GetDebugInfo(const std::string& id) const;