ci: Use testcontainer for running tests

Instead of relying on services to be set up properly, use
`testcontainers` to create a container to run the test inside.

Tests can be decorated with `doctor.unittest.requires` to require
either a `postgres` or `timescale` container. If a Postgres container
is requested, it will be read from the environment variable
`TEST_CONTAINER_POSTGRES` and Timescale containers will be read from
`TEST_CONTAINER_TIMESCALE`.

.github/workflows/tests.yml

@@ -17,18 +17,6 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.8", "3.9", "3.10"]
-    services:
-      timescale:
-        image: timescale/timescaledb:latest-pg15
-        env:
-          POSTGRES_PASSWORD: xyzzy
-        options: >-
-          --health-cmd pg_isready
-          --health-interval 10s
-          --health-timeout 5s
-          --health-retries 5
-        ports:
-          - 5432:5432
 
     timeout-minutes: 25
     steps:
@@ -54,4 +42,6 @@ jobs:
           PGUSER: postgres
           PGPASSWORD: xyzzy
           PGPORT: 5432
+          TEST_CONTAINER_TIMESCALE: timescale/timescaledb:latest-pg15
+          TEST_CONTAINER_POSTGRES: postgres:15
         run: python -m pytest 

src/doctor/rules/__init__.py

@@ -28,7 +28,7 @@ def is_rule_file(fname):
     """Check if file is a rules file."""
     if not isfile(fname):
         return False
-    if fname.endswith(['__init__.py', '_test.py']) or fname.startswith("test_"):
+    if fname.endswith(('__init__.py', '_test.py')) or fname.startswith("test_"):
         return False
     return True
 

src/doctor/rules/compression_test.py

@@ -17,13 +17,15 @@
 import os
 import unittest
 import psycopg2
+import doctor.unittest
 
 from psycopg2.extras import RealDictCursor
 from timescaledb import Hypertable
 
 from doctor.rules.compression import LinearSegmentBy, PointlessSegmentBy
 
-class TestCompressionRules(unittest.TestCase):
+@doctor.unittest.requires('timescaledb')
+class TestCompressionRules(doctor.unittest.TestCase):
     """Test compression rules.
 
     This will create a hypertable where we segment-by a column that
@@ -34,24 +36,15 @@ class TestCompressionRules(unittest.TestCase):
 
     def setUp(self):
         """Set up unit tests for compression rules."""
-        user = os.getenv("PGUSER")
-        host = os.getenv("PGHOST")
-        port = os.getenv("PGPORT") or "5432"
-        dbname = os.getenv("PGDATABASE")
-        password = os.getenv("PGPASSWORD")
-        print(f"connecting to {host}:{port} database {dbname}")
-        self.__conn = psycopg2.connect(dbname=dbname, user=user, host=host,
-                                       password=password, port=port,
-                                       cursor_factory=RealDictCursor)
         table = Hypertable("conditions", "time", {
                            'time': "timestamptz not null",
                            'device_id': "integer",
                            'user_id': "integer",
                            'temperature': "float"
         })
-        table.create(self.__conn)
+        table.create(self.connection)
 
-        with self.__conn.cursor() as cursor:
+        with self.connection.cursor() as cursor:
             cursor.execute(
                 "INSERT INTO conditions "
                 "SELECT time, (random()*30)::int, 1, random()*80 - 40 "
@@ -64,17 +57,13 @@ def setUp(self):
                 ")"
                 )
             cursor.execute("ANALYZE conditions")
-        self.__conn.commit()
+        self.connection.commit()
 
     def tearDown(self):
         """Tear down compression rules test."""
-        with self.__conn.cursor() as cursor:
+        with self.connection.cursor() as cursor:
             cursor.execute("DROP TABLE conditions")
-        self.__conn.commit()
-
-    def run_rule(self, rule):
-        """Run rule and return messages."""
-        return rule.execute(self.__conn, rule.message)
+        self.connection.commit()
 
     def test_segmentby(self):
         """Test rule for detecting bad choice for segment-by column."""

src/doctor/unittest.py

@@ -0,0 +1,80 @@
+# Copyright 2023 Timescale, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests support for Timescale Doctor rules."""
+
+import os
+import unittest
+import psycopg2
+import os
+
+from testcontainers.postgres import PostgresContainer
+from psycopg2.extras import RealDictCursor
+
+def requires(extension):
+    """Decorate test case with required extension."""
+    def set_container(cls):
+        if extension == 'timescaledb':
+            cls.container_name = os.environ.get('TEST_CONTAINER_TIMESCALE')
+        else:
+            cls.container_name = os.environ.get('TEST_CONTAINER_POSTGRES')
+        return cls
+    return set_container
+
+class TestCase(unittest.TestCase):
+    """Base class for Timescale Doctor unit tests.
+
+    Test cases are executed in a container that depends on the
+    requirements of the test case. These are typically used together
+    with the ``required`` decorator to attach an image to the test
+    case and run the tests there.
+
+        import doctor.unittest
+
+        @doctor.unittest.requires('timescaledb')
+        class TestCompressionRules(doctor.unittest.TestCase):
+            ...
+    """
+
+    container_name = None
+
+    @property
+    def connection(self):
+        """Get database connection."""
+        return self.__connection
+
+    @property
+    def container(self):
+        """Get container the test is running in."""
+        return self.__container
+
+    def run_rule(self, rule):
+        """Run rule and return messages."""
+        return rule.execute(self.connection, rule.message)
+
+    @classmethod
+    def setUpClass(self):
+        """Start a container for the test case."""
+        container_name = self.container_name
+        if container_name is None:
+            container_name = os.environ.get('TEST_CONTAINER_POSTGRES', 'postgres:latest')
+        self.__container = PostgresContainer(container_name).start()
+        connstring = self.__container.get_connection_url().replace("+psycopg2", "")
+        self.__connection = psycopg2.connect(connstring, cursor_factory=RealDictCursor)
+
+    @classmethod
+    def tearDownClass(self):
+        """Close the connection and stop the container."""
+        self.__connection.close()
+        self.__container.stop()