feat: add unit tests

Add unit tests and unit test framework for existing functionality.

src/doctor/__init__.py

@@ -71,6 +71,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

@@ -17,8 +17,8 @@
 import argparse
 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

@@ -31,7 +31,7 @@
 
 @doctor.register
 @dataclass
-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,84 @@
+# 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") or os.getlogin())
+        host = os.getenv("PGHOST")
+        port = os.getenv("PGPORT") or "5432"
+        dbname = (os.getenv("PGDATABASE") or os.getlogin())
+        self.__conn = psycopg2.connect(dbname=dbname, user=user, host=host, 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))