should be it

external/duckdb/test/helpers/test_helpers.cpp

@@ -0,0 +1,383 @@
+// #define CATCH_CONFIG_RUNNER
+#include "catch.hpp"
+
+#include "duckdb/common/file_system.hpp"
+#include "duckdb/common/value_operations/value_operations.hpp"
+#include "compare_result.hpp"
+#include "duckdb/main/query_result.hpp"
+#include "test_helpers.hpp"
+#include "duckdb/parser/parsed_data/copy_info.hpp"
+#include "duckdb/main/client_context.hpp"
+#include "duckdb/execution/operator/csv_scanner/string_value_scanner.hpp"
+#include "duckdb/common/case_insensitive_map.hpp"
+#include "test_config.hpp"
+#include "pid.hpp"
+#include "duckdb/function/table/read_csv.hpp"
+#include "duckdb/storage/storage_info.hpp"
+#include <cmath>
+#include <fstream>
+
+using namespace std;
+
+#define TESTING_DIRECTORY_NAME "duckdb_unittest_tempdir"
+
+namespace duckdb {
+static string custom_test_directory;
+static case_insensitive_set_t required_requires;
+static bool delete_test_path = true;
+
+bool NO_FAIL(QueryResult &result) {
+	if (result.HasError()) {
+		fprintf(stderr, "Query failed with message: %s\n", result.GetError().c_str());
+	}
+	return !result.HasError();
+}
+
+bool NO_FAIL(duckdb::unique_ptr<QueryResult> result) {
+	return NO_FAIL(*result);
+}
+
+void TestDeleteDirectory(string path) {
+	duckdb::unique_ptr<FileSystem> fs = FileSystem::CreateLocal();
+	try {
+		if (fs->DirectoryExists(path)) {
+			fs->RemoveDirectory(path);
+		}
+	} catch (...) {
+	}
+}
+
+void TestDeleteFile(string path) {
+	duckdb::unique_ptr<FileSystem> fs = FileSystem::CreateLocal();
+	try {
+		fs->TryRemoveFile(path);
+	} catch (...) {
+	}
+}
+
+void TestChangeDirectory(string path) {
+	// set the base path for the tests
+	FileSystem::SetWorkingDirectory(path);
+}
+
+string TestGetCurrentDirectory() {
+	return FileSystem::GetWorkingDirectory();
+}
+
+void DeleteDatabase(string path) {
+	if (!custom_test_directory.empty()) {
+		return;
+	}
+	TestDeleteFile(path);
+	TestDeleteFile(path + ".wal");
+}
+
+void TestCreateDirectory(string path) {
+	duckdb::unique_ptr<FileSystem> fs = FileSystem::CreateLocal();
+	fs->CreateDirectory(path);
+}
+
+string TestJoinPath(string path1, string path2) {
+	duckdb::unique_ptr<FileSystem> fs = FileSystem::CreateLocal();
+	return fs->JoinPath(path1, path2);
+}
+
+void SetTestDirectory(string path) {
+	custom_test_directory = path;
+}
+
+void AddRequire(string require) {
+	required_requires.insert(require);
+}
+
+bool IsRequired(string require) {
+	return required_requires.count(require);
+}
+
+string GetTestDirectory() {
+	if (custom_test_directory.empty()) {
+		return TESTING_DIRECTORY_NAME;
+	}
+	return custom_test_directory;
+}
+
+string TestDirectoryPath() {
+	duckdb::unique_ptr<FileSystem> fs = FileSystem::CreateLocal();
+	auto test_directory = GetTestDirectory();
+	if (!fs->DirectoryExists(test_directory)) {
+		fs->CreateDirectory(test_directory);
+	}
+	string path;
+	if (custom_test_directory.empty()) {
+		// add the PID to the test directory - but only if it was not specified explicitly by the user
+		auto pid = getpid();
+		path = fs->JoinPath(test_directory, to_string(pid));
+	} else {
+		path = test_directory;
+	}
+	if (!fs->DirectoryExists(path)) {
+		fs->CreateDirectory(path);
+	}
+	return path;
+}
+
+void SetDeleteTestPath(bool delete_path) {
+	delete_test_path = delete_path;
+}
+
+bool DeleteTestPath() {
+	return delete_test_path;
+}
+
+void ClearTestDirectory() {
+	if (!DeleteTestPath()) {
+		return;
+	}
+	duckdb::unique_ptr<FileSystem> fs = FileSystem::CreateLocal();
+	auto test_dir = TestDirectoryPath();
+	// try to clear any files we created in the test directory
+	fs->ListFiles(test_dir, [&](const string &file, bool is_dir) {
+		auto full_path = fs->JoinPath(test_dir, file);
+		try {
+			if (is_dir) {
+				fs->RemoveDirectory(full_path);
+			} else {
+				fs->RemoveFile(full_path);
+			}
+		} catch (...) {
+			// skip
+		}
+	});
+}
+
+string TestCreatePath(string suffix) {
+	return TestJoinPath(TestDirectoryPath(), suffix);
+}
+
+bool TestIsInternalError(unordered_set<string> &internal_error_messages, const string &error) {
+	for (auto &error_message : internal_error_messages) {
+		if (StringUtil::Contains(error, error_message)) {
+			return true;
+		}
+	}
+	return false;
+}
+
+unique_ptr<DBConfig> GetTestConfig() {
+	auto &test_config = TestConfiguration::Get();
+
+	auto result = make_uniq<DBConfig>();
+#ifndef DUCKDB_ALTERNATIVE_VERIFY
+	result->options.checkpoint_wal_size = test_config.GetCheckpointWALSize();
+	result->options.checkpoint_on_shutdown = test_config.GetCheckpointOnShutdown();
+#else
+	result->options.checkpoint_on_shutdown = false;
+#endif
+	result->options.abort_on_wal_failure = true;
+#ifdef DUCKDB_RUN_SLOW_VERIFIERS
+	// This mode isn't slow, but we want test coverage both when it's enabled
+	// and when it's not, so we enable only when DUCKDB_RUN_SLOW_VERIFIERS is set.
+	result->options.trim_free_blocks = true;
+#endif
+	result->options.allow_unsigned_extensions = true;
+	auto storage_version = test_config.GetStorageVersion();
+	if (!storage_version.empty()) {
+		result->options.serialization_compatibility = SerializationCompatibility::FromString(storage_version);
+	}
+
+	auto max_threads = test_config.GetMaxThreads();
+	if (max_threads.IsValid()) {
+		result->options.maximum_threads = max_threads.GetIndex();
+	}
+
+	auto block_alloc_size = test_config.GetBlockAllocSize();
+	if (block_alloc_size.IsValid()) {
+		Storage::VerifyBlockAllocSize(block_alloc_size.GetIndex());
+		result->options.default_block_alloc_size = block_alloc_size.GetIndex();
+	}
+
+	result->options.debug_initialize = test_config.GetDebugInitialize();
+	result->options.set_variables.emplace("debug_verify_vector",
+	                                      EnumUtil::ToString(test_config.GetVectorVerification()));
+	return result;
+}
+
+bool CHECK_COLUMN(QueryResult &result_, size_t column_number, vector<duckdb::Value> values) {
+	if (result_.type == QueryResultType::STREAM_RESULT) {
+		fprintf(stderr, "Unexpected stream query result in CHECK_COLUMN\n");
+		return false;
+	}
+	auto &result = (MaterializedQueryResult &)result_;
+	if (result.HasError()) {
+		fprintf(stderr, "Query failed with message: %s\n", result.GetError().c_str());
+		return false;
+	}
+	if (result.names.size() != result.types.size()) {
+		// column names do not match
+		result.Print();
+		return false;
+	}
+	if (values.empty()) {
+		if (result.RowCount() != 0) {
+			result.Print();
+			return false;
+		} else {
+			return true;
+		}
+	}
+	if (result.RowCount() == 0) {
+		result.Print();
+		return false;
+	}
+	if (column_number >= result.types.size()) {
+		result.Print();
+		return false;
+	}
+	for (idx_t row_idx = 0; row_idx < values.size(); row_idx++) {
+		auto value = result.GetValue(column_number, row_idx);
+		// NULL <> NULL, hence special handling
+		if (value.IsNull() && values[row_idx].IsNull()) {
+			continue;
+		}
+
+		if (!Value::DefaultValuesAreEqual(value, values[row_idx])) {
+			// FAIL("Incorrect result! Got " + vector.GetValue(j).ToString()
+			// +
+			//      " but expected " + values[i + j].ToString());
+			result.Print();
+			return false;
+		}
+	}
+	return true;
+}
+
+bool CHECK_COLUMN(duckdb::unique_ptr<duckdb::QueryResult> &result, size_t column_number, vector<duckdb::Value> values) {
+	if (result->type == QueryResultType::STREAM_RESULT) {
+		auto &stream = (StreamQueryResult &)*result;
+		result = stream.Materialize();
+	}
+	return CHECK_COLUMN(*result, column_number, values);
+}
+
+bool CHECK_COLUMN(duckdb::unique_ptr<duckdb::MaterializedQueryResult> &result, size_t column_number,
+                  vector<duckdb::Value> values) {
+	return CHECK_COLUMN((QueryResult &)*result, column_number, values);
+}
+
+string compare_csv(duckdb::QueryResult &result, string csv, bool header) {
+	D_ASSERT(result.type == QueryResultType::MATERIALIZED_RESULT);
+	auto &materialized = (MaterializedQueryResult &)result;
+	if (materialized.HasError()) {
+		fprintf(stderr, "Query failed with message: %s\n", materialized.GetError().c_str());
+		return materialized.GetError();
+	}
+	string error;
+	if (!compare_result(csv, materialized.Collection(), materialized.types, header, error)) {
+		return error;
+	}
+	return "";
+}
+
+string compare_csv_collection(duckdb::ColumnDataCollection &collection, string csv, bool header) {
+	string error;
+	if (!compare_result(csv, collection, collection.Types(), header, error)) {
+		return error;
+	}
+	return "";
+}
+
+string show_diff(DataChunk &left, DataChunk &right) {
+	if (left.ColumnCount() != right.ColumnCount()) {
+		return StringUtil::Format("Different column counts: %d vs %d", (int)left.ColumnCount(),
+		                          (int)right.ColumnCount());
+	}
+	if (left.size() != right.size()) {
+		return StringUtil::Format("Different sizes: %zu vs %zu", left.size(), right.size());
+	}
+	string difference;
+	for (size_t i = 0; i < left.ColumnCount(); i++) {
+		bool has_differences = false;
+		auto &left_vector = left.data[i];
+		auto &right_vector = right.data[i];
+		string left_column = StringUtil::Format("Result\n------\n%s [", left_vector.GetType().ToString().c_str());
+		string right_column = StringUtil::Format("Expect\n------\n%s [", right_vector.GetType().ToString().c_str());
+		if (left_vector.GetType() == right_vector.GetType()) {
+			for (size_t j = 0; j < left.size(); j++) {
+				auto left_value = left_vector.GetValue(j);
+				auto right_value = right_vector.GetValue(j);
+				if (!Value::DefaultValuesAreEqual(left_value, right_value)) {
+					left_column += left_value.ToString() + ",";
+					right_column += right_value.ToString() + ",";
+					has_differences = true;
+				} else {
+					left_column += "_,";
+					right_column += "_,";
+				}
+			}
+		} else {
+			left_column += "...";
+			right_column += "...";
+		}
+		left_column += "]\n";
+		right_column += "]\n";
+		if (has_differences) {
+			difference += StringUtil::Format("Difference in column %d:\n", i);
+			difference += left_column + "\n" + right_column + "\n";
+		}
+	}
+	return difference;
+}
+
+//! Compares the result of a pipe-delimited CSV with the given DataChunk
+//! Returns true if they are equal, and stores an error_message otherwise
+bool compare_result(string csv, ColumnDataCollection &collection, vector<LogicalType> sql_types, bool has_header,
+                    string &error_message) {
+	D_ASSERT(collection.Count() == 0 || collection.Types().size() == sql_types.size());
+
+	// create the csv on disk
+	auto csv_path = TestCreatePath("__test_csv_path.csv");
+	ofstream f(csv_path);
+	f << csv;
+	f.close();
+
+	// set up the CSV reader
+	CSVReaderOptions options;
+	options.auto_detect = false;
+	options.dialect_options.state_machine_options.delimiter = {"|"};
+	options.dialect_options.header = has_header;
+	options.dialect_options.state_machine_options.quote = '\"';
+	options.dialect_options.state_machine_options.escape = '\"';
+	options.file_path = csv_path;
+	options.dialect_options.num_cols = sql_types.size();
+	// set up the intermediate result chunk
+	DataChunk parsed_result;
+	parsed_result.Initialize(Allocator::DefaultAllocator(), sql_types);
+
+	DuckDB db;
+	Connection con(db);
+	MultiFileOptions file_options;
+	auto scanner_ptr = StringValueScanner::GetCSVScanner(*con.context, options, file_options);
+	auto &scanner = *scanner_ptr;
+	ColumnDataCollection csv_data_collection(*con.context, sql_types);
+	while (!scanner.FinishedIterator()) {
+		// parse a chunk from the CSV file
+		try {
+			parsed_result.Reset();
+			scanner.Flush(parsed_result);
+		} catch (std::exception &ex) {
+			error_message = "Could not parse CSV: " + string(ex.what());
+			return false;
+		}
+		if (parsed_result.size() == 0) {
+			break;
+		}
+		csv_data_collection.Append(parsed_result);
+	}
+	string error;
+	if (!ColumnDataCollection::ResultEquals(collection, csv_data_collection, error_message)) {
+		return false;
+	}
+	return true;
+}
+
+} // namespace duckdb