summaryrefslogtreecommitdiffstats
path: root/tests/helpers.py
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--tests/helpers.py130
1 files changed, 130 insertions, 0 deletions
diff --git a/tests/helpers.py b/tests/helpers.py
new file mode 100644
index 0000000..d4edb14
--- /dev/null
+++ b/tests/helpers.py
@@ -0,0 +1,130 @@
+import os
+
+FILE_DIR = os.path.dirname(__file__)
+FIXTURES_DIR = os.path.join(FILE_DIR, "fixtures")
+
+
+def _filter_comments(s):
+ return "\n".join(
+ [line for line in s.splitlines() if line and not line.startswith("--")]
+ )
+
+
+def _extract_meta(sql):
+ meta = {}
+ sql_lines = sql.split("\n")
+ i = 0
+ while sql_lines[i].startswith("#"):
+ key, val = sql_lines[i].split(":", maxsplit=1)
+ meta[key.lstrip("#").strip()] = val.strip()
+ i += 1
+ sql = "\n".join(sql_lines[i:])
+ return sql, meta
+
+
+def assert_logger_contains(message, logger, level="error"):
+ output = "\n".join(
+ str(args[0][0]) for args in getattr(logger, level).call_args_list
+ )
+ assert message in output
+
+
+def load_sql_fixtures(filename):
+ with open(os.path.join(FIXTURES_DIR, filename), encoding="utf-8") as f:
+ for sql in _filter_comments(f.read()).splitlines():
+ yield sql
+
+
+def load_sql_fixture_pairs(filename):
+ with open(os.path.join(FIXTURES_DIR, filename), encoding="utf-8") as f:
+ statements = _filter_comments(f.read()).split(";")
+
+ size = len(statements)
+
+ for i in range(0, size, 2):
+ if i + 1 < size:
+ sql = statements[i].strip()
+ sql, meta = _extract_meta(sql)
+ expected = statements[i + 1].strip()
+ yield meta, sql, expected
+
+
+TPCH_SCHEMA = {
+ "lineitem": {
+ "l_orderkey": "uint64",
+ "l_partkey": "uint64",
+ "l_suppkey": "uint64",
+ "l_linenumber": "uint64",
+ "l_quantity": "float64",
+ "l_extendedprice": "float64",
+ "l_discount": "float64",
+ "l_tax": "float64",
+ "l_returnflag": "string",
+ "l_linestatus": "string",
+ "l_shipdate": "date32",
+ "l_commitdate": "date32",
+ "l_receiptdate": "date32",
+ "l_shipinstruct": "string",
+ "l_shipmode": "string",
+ "l_comment": "string",
+ },
+ "orders": {
+ "o_orderkey": "uint64",
+ "o_custkey": "uint64",
+ "o_orderstatus": "string",
+ "o_totalprice": "float64",
+ "o_orderdate": "date32",
+ "o_orderpriority": "string",
+ "o_clerk": "string",
+ "o_shippriority": "int32",
+ "o_comment": "string",
+ },
+ "customer": {
+ "c_custkey": "uint64",
+ "c_name": "string",
+ "c_address": "string",
+ "c_nationkey": "uint64",
+ "c_phone": "string",
+ "c_acctbal": "float64",
+ "c_mktsegment": "string",
+ "c_comment": "string",
+ },
+ "part": {
+ "p_partkey": "uint64",
+ "p_name": "string",
+ "p_mfgr": "string",
+ "p_brand": "string",
+ "p_type": "string",
+ "p_size": "int32",
+ "p_container": "string",
+ "p_retailprice": "float64",
+ "p_comment": "string",
+ },
+ "supplier": {
+ "s_suppkey": "uint64",
+ "s_name": "string",
+ "s_address": "string",
+ "s_nationkey": "uint64",
+ "s_phone": "string",
+ "s_acctbal": "float64",
+ "s_comment": "string",
+ },
+ "partsupp": {
+ "ps_partkey": "uint64",
+ "ps_suppkey": "uint64",
+ "ps_availqty": "int32",
+ "ps_supplycost": "float64",
+ "ps_comment": "string",
+ },
+ "nation": {
+ "n_nationkey": "uint64",
+ "n_name": "string",
+ "n_regionkey": "uint64",
+ "n_comment": "string",
+ },
+ "region": {
+ "r_regionkey": "uint64",
+ "r_name": "string",
+ "r_comment": "string",
+ },
+}