gh-108550: Speed up sqlite3 tests (#108551)

erlend-aasland · serhiy-storchaka · web-flow · commit 0e8b3fc718c8 · 2023-08-28T14:17:34.000+02:00
Refactor the CLI so we can easily invoke it and mock command-line
arguments. Adapt the CLI tests so we no longer have to launch a
separate process.

Disable the busy handler for all concurrency tests; we have full
control over the order of the SQLite C API calls, so we can safely
do this.

The sqlite3 test suite now completes ~8 times faster than before.

Co-authored-by: Serhiy Storchaka &lt;storchaka@gmail.com&gt;
diff --git a/Lib/sqlite3/__main__.py b/Lib/sqlite3/__main__.py
@@ -62,7 +62,7 @@ def runsource(self, source, filename="<input>", symbol="single"):
         return False
 
 
-def main():
+def main(*args):
     parser = ArgumentParser(
         description="Python sqlite3 CLI",
         prog="python -m sqlite3",
@@ -86,7 +86,7 @@ def main():
         version=f"SQLite version {sqlite3.sqlite_version}",
         help="Print underlying SQLite library version",
     )
-    args = parser.parse_args()
+    args = parser.parse_args(*args)
 
     if args.filename == ":memory:":
         db_name = "a transient in-memory database"
@@ -120,5 +120,8 @@ def main():
     finally:
         con.close()
 
+    sys.exit(0)
 
-main()
+
+if __name__ == "__main__":
+    main(sys.argv)
diff --git a/Lib/test/test_sqlite3/test_cli.py b/Lib/test/test_sqlite3/test_cli.py
@@ -1,42 +1,35 @@
 """sqlite3 CLI tests."""
-
-import sqlite3 as sqlite
-import subprocess
-import sys
+import sqlite3
 import unittest
 
-from test.support import SHORT_TIMEOUT, requires_subprocess
+from sqlite3.__main__ import main as cli
 from test.support.os_helper import TESTFN, unlink
+from test.support import captured_stdout, captured_stderr, captured_stdin
 
 
-@requires_subprocess()
 class CommandLineInterface(unittest.TestCase):
 
     def _do_test(self, *args, expect_success=True):
-        with subprocess.Popen(
-            [sys.executable, "-Xutf8", "-m", "sqlite3", *args],
-            encoding="utf-8",
-            bufsize=0,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        ) as proc:
-            proc.wait()
-            if expect_success == bool(proc.returncode):
-                self.fail("".join(proc.stderr))
-            stdout = proc.stdout.read()
-            stderr = proc.stderr.read()
-            if expect_success:
-                self.assertEqual(stderr, "")
-            else:
-                self.assertEqual(stdout, "")
-            return stdout, stderr
+        with (
+            captured_stdout() as out,
+            captured_stderr() as err,
+            self.assertRaises(SystemExit) as cm
+        ):
+            cli(args)
+        return out.getvalue(), err.getvalue(), cm.exception.code
 
     def expect_success(self, *args):
-        out, _ = self._do_test(*args)
+        out, err, code = self._do_test(*args)
+        self.assertEqual(code, 0,
+                         "\n".join([f"Unexpected failure: {args=}", out, err]))
+        self.assertEqual(err, "")
         return out
 
     def expect_failure(self, *args):
-        _, err = self._do_test(*args, expect_success=False)
+        out, err, code = self._do_test(*args, expect_success=False)
+        self.assertNotEqual(code, 0,
+                            "\n".join([f"Unexpected failure: {args=}", out, err]))
+        self.assertEqual(out, "")
         return err
 
     def test_cli_help(self):
@@ -45,7 +38,7 @@ def test_cli_help(self):
 
     def test_cli_version(self):
         out = self.expect_success("-v")
-        self.assertIn(sqlite.sqlite_version, out)
+        self.assertIn(sqlite3.sqlite_version, out)
 
     def test_cli_execute_sql(self):
         out = self.expect_success(":memory:", "select 1")
@@ -68,87 +61,68 @@ def test_cli_on_disk_db(self):
         self.assertIn("(0,)", out)
 
 
-@requires_subprocess()
 class InteractiveSession(unittest.TestCase):
-    TIMEOUT = SHORT_TIMEOUT / 10.
     MEMORY_DB_MSG = "Connected to a transient in-memory database"
     PS1 = "sqlite> "
     PS2 = "... "
 
-    def start_cli(self, *args):
-        return subprocess.Popen(
-            [sys.executable, "-Xutf8", "-m", "sqlite3", *args],
-            encoding="utf-8",
-            bufsize=0,
-            stdin=subprocess.PIPE,
-            # Note: the banner is printed to stderr, the prompt to stdout.
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-        )
-
-    def expect_success(self, proc):
-        proc.wait()
-        if proc.returncode:
-            self.fail("".join(proc.stderr))
+    def run_cli(self, *args, commands=()):
+        with (
+            captured_stdin() as stdin,
+            captured_stdout() as stdout,
+            captured_stderr() as stderr,
+            self.assertRaises(SystemExit) as cm
+        ):
+            for cmd in commands:
+                stdin.write(cmd + "\n")
+            stdin.seek(0)
+            cli(args)
+
+        out = stdout.getvalue()
+        err = stderr.getvalue()
+        self.assertEqual(cm.exception.code, 0,
+                         f"Unexpected failure: {args=}\n{out}\n{err}")
+        return out, err
 
     def test_interact(self):
-        with self.start_cli() as proc:
-            out, err = proc.communicate(timeout=self.TIMEOUT)
-            self.assertIn(self.MEMORY_DB_MSG, err)
-            self.assertIn(self.PS1, out)
-            self.expect_success(proc)
+        out, err = self.run_cli()
+        self.assertIn(self.MEMORY_DB_MSG, err)
+        self.assertIn(self.PS1, out)
 
     def test_interact_quit(self):
-        with self.start_cli() as proc:
-            out, err = proc.communicate(input=".quit", timeout=self.TIMEOUT)
-            self.assertIn(self.MEMORY_DB_MSG, err)
-            self.assertIn(self.PS1, out)
-            self.expect_success(proc)
+        out, err = self.run_cli(commands=(".quit",))
+        self.assertIn(self.PS1, out)
 
     def test_interact_version(self):
-        with self.start_cli() as proc:
-            out, err = proc.communicate(input=".version", timeout=self.TIMEOUT)
-            self.assertIn(self.MEMORY_DB_MSG, err)
-            self.assertIn(sqlite.sqlite_version, out)
-            self.expect_success(proc)
+        out, err = self.run_cli(commands=(".version",))
+        self.assertIn(self.MEMORY_DB_MSG, err)
+        self.assertIn(sqlite3.sqlite_version, out)
 
     def test_interact_valid_sql(self):
-        with self.start_cli() as proc:
-            out, err = proc.communicate(input="select 1;",
-                                        timeout=self.TIMEOUT)
-            self.assertIn(self.MEMORY_DB_MSG, err)
-            self.assertIn("(1,)", out)
-            self.expect_success(proc)
+        out, err = self.run_cli(commands=("SELECT 1;",))
+        self.assertIn(self.MEMORY_DB_MSG, err)
+        self.assertIn("(1,)", out)
 
     def test_interact_valid_multiline_sql(self):
-        with self.start_cli() as proc:
-            out, err = proc.communicate(input="select 1\n;",
-                                        timeout=self.TIMEOUT)
-            self.assertIn(self.MEMORY_DB_MSG, err)
-            self.assertIn(self.PS2, out)
-            self.assertIn("(1,)", out)
-            self.expect_success(proc)
+        out, err = self.run_cli(commands=("SELECT 1\n;",))
+        self.assertIn(self.MEMORY_DB_MSG, err)
+        self.assertIn(self.PS2, out)
+        self.assertIn("(1,)", out)
 
     def test_interact_invalid_sql(self):
-        with self.start_cli() as proc:
-            out, err = proc.communicate(input="sel;", timeout=self.TIMEOUT)
-            self.assertIn(self.MEMORY_DB_MSG, err)
-            self.assertIn("OperationalError (SQLITE_ERROR)", err)
-            self.expect_success(proc)
+        out, err = self.run_cli(commands=("sel;",))
+        self.assertIn(self.MEMORY_DB_MSG, err)
+        self.assertIn("OperationalError (SQLITE_ERROR)", err)
 
     def test_interact_on_disk_file(self):
         self.addCleanup(unlink, TESTFN)
-        with self.start_cli(TESTFN) as proc:
-            out, err = proc.communicate(input="create table t(t);",
-                                        timeout=self.TIMEOUT)
-            self.assertIn(TESTFN, err)
-            self.assertIn(self.PS1, out)
-            self.expect_success(proc)
-        with self.start_cli(TESTFN, "select count(t) from t") as proc:
-            out = proc.stdout.read()
-            err = proc.stderr.read()
-            self.assertIn("(0,)", out)
-            self.expect_success(proc)
+
+        out, err = self.run_cli(TESTFN, commands=("CREATE TABLE t(t);",))
+        self.assertIn(TESTFN, err)
+        self.assertIn(self.PS1, out)
+
+        out, _ = self.run_cli(TESTFN, commands=("SELECT count(t) FROM t;",))
+        self.assertIn("(0,)", out)
 
 
 if __name__ == "__main__":
diff --git a/Lib/test/test_sqlite3/test_dbapi.py b/Lib/test/test_sqlite3/test_dbapi.py
@@ -1854,7 +1854,7 @@ def test_on_conflict_replace(self):
 
 @requires_subprocess()
 class MultiprocessTests(unittest.TestCase):
-    CONNECTION_TIMEOUT = SHORT_TIMEOUT / 1000.  # Defaults to 30 ms
+    CONNECTION_TIMEOUT = 0  # Disable the busy timeout.
 
     def tearDown(self):
         unlink(TESTFN)
diff --git a/Lib/test/test_sqlite3/test_transactions.py b/Lib/test/test_sqlite3/test_transactions.py
@@ -24,23 +24,21 @@
 import sqlite3 as sqlite
 from contextlib import contextmanager
 
-from test.support import LOOPBACK_TIMEOUT
 from test.support.os_helper import TESTFN, unlink
 from test.support.script_helper import assert_python_ok
 
 from .util import memory_database
 from .util import MemoryDatabaseMixin
 
 
-TIMEOUT = LOOPBACK_TIMEOUT / 10
-
-
 class TransactionTests(unittest.TestCase):
     def setUp(self):
-        self.con1 = sqlite.connect(TESTFN, timeout=TIMEOUT)
+        # We can disable the busy handlers, since we control
+        # the order of SQLite C API operations.
+        self.con1 = sqlite.connect(TESTFN, timeout=0)
         self.cur1 = self.con1.cursor()
 
-        self.con2 = sqlite.connect(TESTFN, timeout=TIMEOUT)
+        self.con2 = sqlite.connect(TESTFN, timeout=0)
         self.cur2 = self.con2.cursor()
 
     def tearDown(self):
@@ -120,10 +118,8 @@ def test_raise_timeout(self):
             self.cur2.execute("insert into test(i) values (5)")
 
     def test_locking(self):
-        """
-        This tests the improved concurrency with pysqlite 2.3.4. You needed
-        to roll back con2 before you could commit con1.
-        """
+        # This tests the improved concurrency with pysqlite 2.3.4. You needed
+        # to roll back con2 before you could commit con1.
         self.cur1.execute("create table test(i)")
         self.cur1.execute("insert into test(i) values (5)")
         with self.assertRaises(sqlite.OperationalError):
