Move CSV->Parquet conversion to test register instance (#1656)

kuzudb · Jun 10, 2023 · 540ee67 · 540ee67
1 parent 43c7026
commit 540ee67
Showing 1 changed file with 29 additions and 23 deletions.
diff --git a/test/runner/e2e_test.cpp b/test/runner/e2e_test.cpp
@@ -7,30 +7,43 @@ using ::testing::Test;
 using namespace kuzu::testing;
 using namespace kuzu::common;
 
-class EndToEndTest : public DBTest {
+class EndToEndEnvironment : public testing::Environment {
 public:
-    explicit EndToEndTest(std::string dataset, uint64_t bufferPoolSize,
-        std::vector<std::unique_ptr<TestStatement>> testStatements)
-        : dataset{dataset}, bufferPoolSize{bufferPoolSize}, testStatements{
-                                                                std::move(testStatements)} {}
+    virtual void SetUp() {
+        FileUtils::createDirIfNotExists(
+            TestHelper::appendKuzuRootPath(TestHelper::PARQUET_TEMP_DATASET_PATH));
+    }
+    virtual void TearDown() {
+        FileUtils::removeDir(TestHelper::appendKuzuRootPath(TestHelper::PARQUET_TEMP_DATASET_PATH));
+    }
+};
 
+class EndToEndTest : public DBTest {
+public:
+    explicit EndToEndTest(TestGroup::DatasetType datasetType, std::string dataset,
+        uint64_t bufferPoolSize, std::vector<std::unique_ptr<TestStatement>> testStatements)
+        : datasetType{datasetType}, dataset{dataset}, bufferPoolSize{bufferPoolSize},
+          testStatements{std::move(testStatements)} {}
     void SetUp() override {
+        setUpDataset();
         BaseGraphTest::SetUp();
         systemConfig->bufferPoolSize = bufferPoolSize;
         createDBAndConn();
         initGraph();
     }
-
-    void TearDown() override {
-        FileUtils::removeDir(TestHelper::appendKuzuRootPath(TestHelper::PARQUET_TEMP_DATASET_PATH));
-        FileUtils::removeDir(TestHelper::getTmpTestDir());
+    void setUpDataset() {
+        if (datasetType == TestGroup::DatasetType::CSV_TO_PARQUET) {
+            CSVToParquetConverter::convertCSVDatasetToParquet(dataset);
+        } else {
+            dataset = TestHelper::appendKuzuRootPath("dataset/" + dataset);
+        }
     }
-
+    void TearDown() override { FileUtils::removeDir(TestHelper::getTmpTestDir()); }
     std::string getInputDir() override { return dataset + "/"; }
-
     void TestBody() override { runTest(testStatements); }
 
 private:
+    TestGroup::DatasetType datasetType;
     std::string dataset;
     uint64_t bufferPoolSize;
     std::vector<std::unique_ptr<TestStatement>> testStatements;
@@ -40,24 +53,21 @@ void parseAndRegisterTestGroup(const std::string& path, bool generateTestList =
     auto testParser = std::make_unique<TestParser>(path);
     auto testGroup = std::move(testParser->parseTestFile());
     if (testGroup->isValid() && testGroup->hasStatements()) {
+        auto datasetType = testGroup->datasetType;
         auto dataset = testGroup->dataset;
         auto testCases = std::move(testGroup->testCases);
         auto bufferPoolSize = testGroup->bufferPoolSize;
-        if (testGroup->datasetType == TestGroup::DatasetType::CSV_TO_PARQUET) {
-            CSVToParquetConverter::convertCSVDatasetToParquet(dataset);
-        } else {
-            dataset = TestHelper::appendKuzuRootPath("dataset/" + dataset);
-        }
         for (auto& [testCaseName, testStatements] : testCases) {
             if (generateTestList) {
                 std::ofstream testList(TestHelper::getTestListFile(), std::ios_base::app);
                 testList << testGroup->group + "." + testCaseName + " " + path + "\n";
             }
             testing::RegisterTest(testGroup->group.c_str(), testCaseName.c_str(), nullptr, nullptr,
                 __FILE__, __LINE__,
-                [dataset, bufferPoolSize,
+                [datasetType, dataset, bufferPoolSize,
                     testStatements = std::move(testStatements)]() mutable -> DBTest* {
-                    return new EndToEndTest(dataset, bufferPoolSize, std::move(testStatements));
+                    return new EndToEndTest(
+                        datasetType, dataset, bufferPoolSize, std::move(testStatements));
                 });
         }
     } else {
@@ -69,8 +79,6 @@ void scanTestFiles(const std::string& path) {
     std::string testListFile = TestHelper::appendKuzuRootPath(
         FileUtils::joinPath(TestHelper::E2E_TEST_FILES_DIRECTORY, "test_list"));
     FileUtils::removeFileIfExists(testListFile);
-    FileUtils::createDirIfNotExists(
-        TestHelper::appendKuzuRootPath(TestHelper::PARQUET_TEMP_DATASET_PATH));
     if (std::filesystem::is_regular_file(path)) {
         parseAndRegisterTestGroup(path);
         return;
@@ -97,13 +105,10 @@ std::string findTestFile(std::string testCase) {
 void checkCtestParams(int argc, char** argv) {
     if (argc > 1) {
         std::string argument = argv[1];
-        bool runFromCTest = false;
         if (argument == "--gtest_list_tests") {
             scanTestFiles(TestHelper::appendKuzuRootPath(TestHelper::E2E_TEST_FILES_DIRECTORY));
         }
         if (argument.starts_with("--gtest_filter=")) {
-            FileUtils::createDirIfNotExists(
-                TestHelper::appendKuzuRootPath(TestHelper::PARQUET_TEMP_DATASET_PATH));
             std::string testCaseFile = findTestFile(argument.substr(15));
             if (testCaseFile.empty()) {
                 scanTestFiles(TestHelper::appendKuzuRootPath(TestHelper::E2E_TEST_FILES_DIRECTORY));
@@ -117,6 +122,7 @@ void checkCtestParams(int argc, char** argv) {
 int main(int argc, char** argv) {
     checkCtestParams(argc, argv);
     testing::InitGoogleTest(&argc, argv);
+    testing::AddGlobalTestEnvironment(new EndToEndEnvironment);
     if (argc > 1) {
         auto path = TestHelper::appendKuzuRootPath(
             FileUtils::joinPath(TestHelper::E2E_TEST_FILES_DIRECTORY, argv[1]));