feat: add unit tests

.github/workflows/tests.yml

@@ -0,0 +1,57 @@
+# This file contains checks to run unit tests.
+#
+# Special thank you to pylint-dev project at GitHub where inspiration
+# was taken.
+name: unit tests
+
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+
+jobs:
+  tests-linux:
+    name: run / ${{ matrix.python-version }} / Linux
+    runs-on: ubuntu-latest
+    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:
+      - uses: actions/[email protected]
+      - name: Set up Python ${{ matrix.python-version }}
+        id: python
+        uses: actions/[email protected]
+        with:
+          python-version: ${{ matrix.python-version }}
+          check-latest: true
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pytest
+          if [ -f requirements.txt ]; then
+            pip install -r requirements.txt
+          else
+            echo "No requirements.txt file found"
+          fi
+      - name: Run PyTest
+        env:
+          PGHOST: localhost
+          PGUSER: postgres
+          PGPASSWORD: xyzzy
+          PGPORT: 5432
+        run: python -m pytest 

requirements.txt

@@ -0,0 +1 @@
+psycopg2>=2.9.2

src/doctor/__init__.py

@@ -89,6 +89,7 @@ def register(cls):
         raise NameError('Class did not define a message', name='message')
     category = cls.__module__.rpartition('.')[2]
     RULES.setdefault(category, {})[cls.__name__] = cls
+    return cls
 
 
 def check_rules(dbname, user, host, port):

src/doctor/cli.py

@@ -18,8 +18,8 @@
 import getpass
 import os
 
-from . import check_rules
-from .rules import load_rules
+from doctor import check_rules
+from doctor.rules import load_rules
 
 
 def parse_arguments():

src/doctor/rules/__init__.py

@@ -24,12 +24,18 @@
 
 from os.path import dirname, basename, isfile, join
 
+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_"):
+        return False
+    return True
+
 def load_rules():
     """Load all rules from package."""
     # Add all files in this directory to be loaded, except __init__.py
-    modules = [
-        basename(f)[:-3] for f in glob.glob(join(dirname(__file__), "*.py"))
-        if isfile(f) and not f.endswith('__init__.py')
-    ]
+    pyfiles = glob.glob(join(dirname(__file__), "*.py"))
+    modules = [basename(f)[:-3] for f in pyfiles if is_rule_file(f)]
     for module in modules:
         importlib.import_module(f".{module}", __name__)

src/doctor/rules/compression.py

@@ -28,7 +28,7 @@
 """
 
 @doctor.register
-class LinearSegmentby(doctor.Rule):
+class LinearSegmentBy(doctor.Rule):
     """Detect segmentby column for compressed table."""
 
     query: str = LINEAR_QUERY

src/doctor/rules/compression_test.py

@@ -0,0 +1,87 @@
+# 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 for compressed hypertable rules."""
+
+import os
+import unittest
+import psycopg2
+
+from psycopg2.extras import RealDictCursor
+from timescaledb import Hypertable
+
+from doctor.rules.compression import LinearSegmentBy, PointlessSegmentBy
+
+class TestCompressionRules(unittest.TestCase):
+    """Test compression rules.
+
+    This will create a hypertable where we segment-by a column that
+    has unique values (the time column) and a column that has only a
+    single value (user_id).
+
+    """
+
+    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)
+
+        with self.__conn.cursor() as cursor:
+            cursor.execute(
+                "INSERT INTO conditions "
+                "SELECT time, (random()*30)::int, 1, random()*80 - 40 "
+                "FROM generate_series(NOW() - INTERVAL '6 days', NOW(), '1 minute') AS time"
+            )
+            cursor.execute(
+                "ALTER TABLE conditions SET ("
+                "    timescaledb.compress,"
+                "    timescaledb.compress_segmentby = 'time, user_id'"
+                ")"
+                )
+            cursor.execute("ANALYZE conditions")
+        self.__conn.commit()
+
+    def tearDown(self):
+        """Tear down compression rules test."""
+        with self.__conn.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)
+
+    def test_segmentby(self):
+        """Test rule for detecting bad choice for segment-by column."""
+        messages = []
+        messages.extend(self.run_rule(LinearSegmentBy()))
+        messages.extend(self.run_rule(PointlessSegmentBy()))
+        self.assertIn(LinearSegmentBy.message.format(attname="time", relation="conditions"),
+                      messages)
+        self.assertIn(PointlessSegmentBy.message.format(attname="user_id", relation="conditions"),
+                      messages)

src/timescaledb/__init__.py

@@ -0,0 +1,52 @@
+# 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.
+
+"""Utility functions for interacting with TimescaleDB."""
+
+from abc import ABC, abstractmethod
+
+# pylint: disable-next=too-few-public-methods
+class SQL(ABC):
+    """Superclass for objects in an SQL database."""
+
+    @abstractmethod
+    def create(self, conn):
+        """Create the SQL object in the database."""
+
+# pylint: disable-next=too-few-public-methods
+class Table(SQL):
+    """Class for a normal PostgreSQL table."""
+
+    def __init__(self, name, columns):
+        self.name = name
+        self.columns = columns
+
+    def create(self, conn):
+        with conn.cursor() as cursor:
+            coldefs = ",".join(f"{name} {defn}" for name, defn in self.columns.items())
+            cursor.execute(f"CREATE TABLE {self.name} ({coldefs})")
+
+# pylint: disable-next=too-few-public-methods
+class Hypertable(Table):
+    """Class for a TimescaleDB hypertable."""
+
+    def __init__(self, name, partcol, columns):
+        super().__init__(name, columns)
+        self.partcol = partcol
+
+    def create(self, conn):
+        super().create(conn)
+        with conn.cursor() as cursor:
+            cursor.execute("SELECT * FROM create_hypertable(%s, %s)",
+                           (self.name, self.partcol))