Merge pull request #2029 from hououou/master

Transaction test rewrite

test/graph_test/graph_test.cpp

@@ -158,6 +158,28 @@ void TestHelper::executeScript(const std::string& cypherScript, Connection& conn
         }
     }
 }
+void BaseGraphTest::createDB(uint64_t checkpointWaitTimeout) {
+    if (database != nullptr) {
+        database.reset();
+    }
+    database = std::make_unique<main::Database>(databasePath, *systemConfig);
+    getTransactionManager(*database)->setCheckPointWaitTimeoutForTransactionsToLeaveInMicros(
+        checkpointWaitTimeout /* 10ms */);
+    spdlog::set_level(spdlog::level::info);
+}
+
+void BaseGraphTest::createConns(const std::set<std::string>& connNames) {
+    if (connNames.size() == 0) { // impart a default connName
+        conn = std::make_unique<main::Connection>(database.get());
+    } else {
+        for (auto connName : connNames) {
+            if (connMap[connName] != nullptr) {
+                connMap[connName].reset();
+            }
+            connMap[connName] = std::make_unique<main::Connection>(database.get());
+        }
+    }
+}
 
 void BaseGraphTest::createDBAndConn() {
     if (database != nullptr) {
@@ -169,8 +191,16 @@ void BaseGraphTest::createDBAndConn() {
 }
 
 void BaseGraphTest::initGraph() {
-    TestHelper::executeScript(getInputDir() + TestHelper::SCHEMA_FILE_NAME, *conn);
-    TestHelper::executeScript(getInputDir() + TestHelper::COPY_FILE_NAME, *conn);
+    if (conn) { // normal conn
+        TestHelper::executeScript(getInputDir() + TestHelper::SCHEMA_FILE_NAME, *conn);
+        TestHelper::executeScript(getInputDir() + TestHelper::COPY_FILE_NAME, *conn);
+    } else {
+        // choose a conn from connMap
+        TestHelper::executeScript(
+            getInputDir() + TestHelper::SCHEMA_FILE_NAME, *(connMap.begin()->second));
+        TestHelper::executeScript(
+            getInputDir() + TestHelper::COPY_FILE_NAME, *(connMap.begin()->second));
+    }
 }
 
 void BaseGraphTest::commitOrRollbackConnection(

test/include/graph_test/graph_test.h

@@ -39,6 +39,10 @@ class BaseGraphTest : public Test {
 
     void createDBAndConn();
 
+    // multiple conns test
+    void createDB(uint64_t checkpointWaitTimeout);
+    void createConns(const std::set<std::string>& connNames);
+
     void initGraph();
 
     void commitOrRollbackConnection(bool isCommit, TransactionTestType transactionTestType) const;
@@ -140,7 +144,10 @@ class BaseGraphTest : public Test {
     std::string databasePath;
     std::unique_ptr<main::SystemConfig> systemConfig;
     std::unique_ptr<main::Database> database;
+    // for normal conn; if it is null, respresents multiple conns, need to use connMap
     std::unique_ptr<main::Connection> conn;
+    // for multiple conns
+    std::unordered_map<std::string, std::unique_ptr<main::Connection>> connMap;
 };
 
 // This class starts database without initializing graph.
@@ -157,10 +164,24 @@ class DBTest : public BaseGraphTest {
         initGraph();
     }
 
-    inline void runTest(const std::vector<std::unique_ptr<TestStatement>>& statements) {
-        TestRunner::runTest(statements, *conn, databasePath);
+    inline void runTest(const std::vector<std::unique_ptr<TestStatement>>& statements,
+        uint64_t checkpointWaitTimeout =
+            common::DEFAULT_CHECKPOINT_WAIT_TIMEOUT_FOR_TRANSACTIONS_TO_LEAVE_IN_MICROS,
+        std::set<std::string> connNames = std::set<std::string>()) {
+        for (auto& statement : statements) {
+            if (statement->reLoadDBFlag) {
+                createDB(checkpointWaitTimeout);
+                createConns(connNames);
+                continue;
+            }
+            if (conn) {
+                TestRunner::runTest(statement.get(), *conn, databasePath);
+            } else {
+                auto conn_name = *statement->conn_name;
+                TestRunner::runTest(statement.get(), *connMap[conn_name], databasePath);
+            }
+        }
     }
 };
-
 } // namespace testing
 } // namespace kuzu

test/include/test_runner/test_group.h

@@ -1,4 +1,6 @@
 #pragma once
+#include <optional>
+#include <set>
 
 #include "main/kuzu.h"
 
@@ -18,6 +20,9 @@ struct TestStatement {
     bool enumerate = false;
     bool checkOutputOrder = false;
     std::string expectedTuplesCSVFile;
+    // for multiple conns
+    std::optional<std::string> conn_name; // default conntion name conn0
+    bool reLoadDBFlag = false;
 };
 
 // Test group is a collection of test cases in a single file.
@@ -31,6 +36,12 @@ struct TestGroup {
     std::unordered_map<std::string, std::vector<std::unique_ptr<TestStatement>>>
         testCasesStatementBlocks;
     uint64_t bufferPoolSize = common::BufferPoolConstants::DEFAULT_BUFFER_POOL_SIZE_FOR_TESTING;
+    // for multiple connections
+    uint64_t checkpointWaitTimeout =
+        common::DEFAULT_CHECKPOINT_WAIT_TIMEOUT_FOR_TRANSACTIONS_TO_LEAVE_IN_MICROS;
+    ;
+    bool multipleConns = false;
+    std::unordered_map<std::string, std::set<std::string>> testCasesConnNames;
 
     enum class DatasetType { CSV, PARQUET, NPY, CSV_TO_PARQUET };
     DatasetType datasetType;

test/include/test_runner/test_parser.h

@@ -34,7 +34,11 @@ enum class TokenType {
     ROLLBACK,
     SEPARATOR,
     STATEMENT,
-    _SKIP_LINE
+    _SKIP_LINE,
+
+    CHECKPOINT_WAIT_TIMEOUT,
+    CREATE_CONNECTION,
+    RELOADDB
 };
 
 const std::unordered_map<std::string, TokenType> tokenMap = {{"-GROUP", TokenType::GROUP},
@@ -48,6 +52,8 @@ const std::unordered_map<std::string, TokenType> tokenMap = {{"-GROUP", TokenTyp
     {"-DEFINE", TokenType::DEFINE}, {"-STATEMENT", TokenType::STATEMENT},
     {"-INSERT_STATEMENT_BLOCK", TokenType::INSERT_STATEMENT_BLOCK},
     {"-ROLLBACK", TokenType::ROLLBACK}, {"-BUFFER_POOL_SIZE", TokenType::BUFFER_POOL_SIZE},
+    {"-CHECKPOINT_WAIT_TIMEOUT", TokenType::CHECKPOINT_WAIT_TIMEOUT},
+    {"-RELOADDB", TokenType::RELOADDB}, {"-CREATE_CONNECTION", TokenType::CREATE_CONNECTION},
     {"]", TokenType::END_OF_STATEMENT_BLOCK}, {"----", TokenType::RESULT},
     {"--", TokenType::SEPARATOR}, {"#", TokenType::EMPTY}};
 
@@ -85,6 +91,7 @@ class TestParser {
     void extractDataset();
     void addStatementBlock(const std::string& blockName, const std::string& testGroupName);
     void replaceVariables(std::string& str);
+    bool extractConnName(std::string& query, TestStatement* statement);
 
     inline bool endOfFile() { return fileStream.eof(); }
     inline void setCursorToPreviousLine() { fileStream.seekg(previousFilePosition); }
@@ -110,7 +117,8 @@ class TestParser {
             });
     }
 
-    TestStatement* extractStatement(TestStatement* currentStatement);
+    TestStatement* extractStatement(
+        TestStatement* currentStatement, const std::string& testCaseName);
     TestStatement* addNewStatement(std::string& name);
 
     // Any value here will be replaced inside the .test files

test/include/test_runner/test_runner.h

@@ -10,8 +10,8 @@ namespace testing {
 
 class TestRunner {
 public:
-    static void runTest(const std::vector<std::unique_ptr<TestStatement>>& statements,
-        main::Connection& conn, std::string& databasePath);
+    static void runTest(
+        TestStatement* statement, main::Connection& conn, std::string& databasePath);
 
     static std::unique_ptr<planner::LogicalPlan> getLogicalPlan(
         const std::string& query, main::Connection& conn);

test/runner/e2e_test.cpp

@@ -10,15 +10,19 @@ using namespace kuzu::common;
 class EndToEndTest : public DBTest {
 public:
     explicit EndToEndTest(TestGroup::DatasetType datasetType, std::string dataset,
-        uint64_t bufferPoolSize, std::vector<std::unique_ptr<TestStatement>> testStatements)
+        uint64_t bufferPoolSize, uint64_t checkpointWaitTimeout,
+        const std::set<std::string>& connNames,
+        std::vector<std::unique_ptr<TestStatement>> testStatements)
         : datasetType{datasetType}, dataset{std::move(dataset)}, bufferPoolSize{bufferPoolSize},
+          checkpointWaitTimeout{checkpointWaitTimeout}, connNames{connNames},
           testStatements{std::move(testStatements)} {}
 
     void SetUp() override {
         setUpDataset();
         BaseGraphTest::SetUp();
         systemConfig->bufferPoolSize = bufferPoolSize;
-        createDBAndConn();
+        createDB(checkpointWaitTimeout);
+        createConns(connNames);
         if (dataset != "empty") {
             initGraph();
         }
@@ -38,15 +42,17 @@ class EndToEndTest : public DBTest {
         FileUtils::removeDir(parquetTempDatasetPath);
     }
 
-    void TestBody() override { runTest(testStatements); }
+    void TestBody() override { runTest(testStatements, checkpointWaitTimeout, connNames); }
     std::string getInputDir() override { return dataset + "/"; }
 
 private:
     TestGroup::DatasetType datasetType;
     std::string dataset;
     std::string parquetTempDatasetPath;
     uint64_t bufferPoolSize;
+    uint64_t checkpointWaitTimeout;
     std::vector<std::unique_ptr<TestStatement>> testStatements;
+    std::set<std::string> connNames;
 
     std::string generateParquetTempDatasetPath() {
         return TestHelper::appendKuzuRootPath(
@@ -64,6 +70,7 @@ void parseAndRegisterTestGroup(const std::string& path, bool generateTestList =
         auto dataset = testGroup->dataset;
         auto testCases = std::move(testGroup->testCases);
         auto bufferPoolSize = testGroup->bufferPoolSize;
+        auto checkpointWaitTimeout = testGroup->checkpointWaitTimeout;
         for (auto& [testCaseName, testStatements] : testCases) {
             if (generateTestList) {
                 std::ofstream testList(TestHelper::getTestListFile(), std::ios_base::app);
@@ -72,12 +79,13 @@ void parseAndRegisterTestGroup(const std::string& path, bool generateTestList =
             if (empty(testCaseName)) {
                 throw TestException("Missing test case name (-CASE) [" + path + "].");
             }
+            auto connNames = testGroup->testCasesConnNames[testCaseName];
             testing::RegisterTest(testGroup->group.c_str(), testCaseName.c_str(), nullptr, nullptr,
                 __FILE__, __LINE__,
-                [datasetType, dataset, bufferPoolSize,
+                [datasetType, dataset, bufferPoolSize, checkpointWaitTimeout, connNames,
                     testStatements = std::move(testStatements)]() mutable -> DBTest* {
-                    return new EndToEndTest(
-                        datasetType, dataset, bufferPoolSize, std::move(testStatements));
+                    return new EndToEndTest(datasetType, dataset, bufferPoolSize,
+                        checkpointWaitTimeout, connNames, std::move(testStatements));
                 });
         }
     } else {