From 6fa7bc3d0b86b21d6674a3211e2b1276a354270e Mon Sep 17 00:00:00 2001
From: Colin Dellow <cldellow@gmail.com>
Date: Sat, 24 Mar 2018 11:27:06 -0400
Subject: [PATCH] Add harness for low memory testing

---
 .gitignore            |  3 +++
 tests/test-failmalloc | 41 +++++++++++++++++++++++++++++++++++++++++
 tests/test-random     |  7 ++++---
 3 files changed, 48 insertions(+), 3 deletions(-)
 create mode 100755 tests/test-failmalloc

diff --git a/.gitignore b/.gitignore
index f22cd63..5868472 100644
--- a/.gitignore
+++ b/.gitignore
@@ -48,3 +48,6 @@
 tests/queries
 /testcase-bootstrap.sql
 /test.db
+/tests/test.db
+/tests/results.bad_alloc
+/tests/libfailmalloc
diff --git a/tests/test-failmalloc b/tests/test-failmalloc
new file mode 100755
index 0000000..2fcb21d
--- /dev/null
+++ b/tests/test-failmalloc
@@ -0,0 +1,41 @@
+#!/bin/bash
+set -euo pipefail
+
+# A harness that runs SQLite with the parquet extension in an environment where malloc randomly
+# fails. "Success" is if the logs don't have any C++ exceptions that talk about std::bad_alloc
+
+ensure_failmalloc() {
+  if [ ! -d libfailmalloc ]; then
+    git clone https://github.com/cldellow/libfailmalloc.git
+  fi
+
+  if [ ! -e libfailmalloc/.libs/libfailmalloc.so ]; then
+    cd libfailmalloc
+    ./configure
+    make
+  fi
+}
+
+run_under_low_memory() {
+  start=$(date +%s%3N)
+  set +e
+  env LD_PRELOAD="$here"/libfailmalloc/.libs/libfailmalloc.so FAILMALLOC_PROBABILITY=0.00001 ./test-random &> results.bad_alloc
+  set -e
+  now=$(date +%s%3N)
+  echo "Bailed after $((now-start)) ms"
+  ! grep std::bad_alloc results.bad_alloc
+}
+
+main() {
+  here=$(dirname "${BASH_SOURCE[0]}")
+  here=$(readlink -f "$here")
+  cd "$here"
+
+  ensure_failmalloc
+  # Sometimes we'll exit due to a Python memory issue, so try a few times.
+  for i in {0..10}; do
+    run_under_low_memory
+  done
+}
+
+main "$@"
diff --git a/tests/test-random b/tests/test-random
index 7907f15..f4fdec0 100755
--- a/tests/test-random
+++ b/tests/test-random
@@ -87,14 +87,15 @@ def test_table(conn, table):
     for i in range(1000):
         test_statement(conn, table, column_values, all_values)
 
-def test_db(db_file, tables):
+def test_db(db_file, extension_file, tables):
     conn = sqlite3.connect(db_file)
     conn.enable_load_extension(True)
-    conn.load_extension('parquet/libparquet.so')
+    conn.load_extension(extension_file)
     conn.enable_load_extension(False)
     for table in tables:
         test_table(conn, table)
 
 if __name__ == '__main__':
     db_file = os.path.abspath(os.path.join(__file__, '..', '..', 'test.db'))
-    test_db(db_file, ['nulls', 'no_nulls'])
+    extension_file = os.path.abspath(os.path.join(__file__, '..', '..', 'parquet', 'libparquet.so'))
+    test_db(db_file, extension_file, ['nulls', 'no_nulls'])