From 611e5edd4439ccd0923bc8493370b66691ab4218 Mon Sep 17 00:00:00 2001 From: Mats Kindahl Date: Thu, 14 Sep 2023 15:11:32 +0200 Subject: [PATCH 1/2] feat: add unit tests Add unit tests and unit test framework for existing functionality. --- src/doctor/__init__.py | 1 + src/doctor/cli.py | 4 +- src/doctor/rules/__init__.py | 14 +++-- src/doctor/rules/compression.py | 2 +- src/doctor/rules/compression_test.py | 84 ++++++++++++++++++++++++++++ src/timescaledb/__init__.py | 52 +++++++++++++++++ 6 files changed, 150 insertions(+), 7 deletions(-) create mode 100644 src/doctor/rules/compression_test.py create mode 100644 src/timescaledb/__init__.py diff --git a/src/doctor/__init__.py b/src/doctor/__init__.py index 8de5707..5d96c64 100644 --- a/src/doctor/__init__.py +++ b/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): diff --git a/src/doctor/cli.py b/src/doctor/cli.py index c16f188..73d2d3b 100644 --- a/src/doctor/cli.py +++ b/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(): diff --git a/src/doctor/rules/__init__.py b/src/doctor/rules/__init__.py index 646c6a9..f82a58d 100644 --- a/src/doctor/rules/__init__.py +++ b/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__) diff --git a/src/doctor/rules/compression.py b/src/doctor/rules/compression.py index 8eb0bcf..f283d03 100644 --- a/src/doctor/rules/compression.py +++ b/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 diff --git a/src/doctor/rules/compression_test.py b/src/doctor/rules/compression_test.py new file mode 100644 index 0000000..be28b43 --- /dev/null +++ b/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) diff --git a/src/timescaledb/__init__.py b/src/timescaledb/__init__.py new file mode 100644 index 0000000..f1597d7 --- /dev/null +++ b/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)) From 42674a6de088de12c4a329d13aaff6c5e908f947 Mon Sep 17 00:00:00 2001 From: Mats Kindahl Date: Thu, 14 Sep 2023 18:34:31 +0200 Subject: [PATCH 2/2] ci: add workflow to run PyTest --- .github/workflows/tests.yml | 57 ++++++++++++++++++++++++++++ requirements.txt | 1 + src/doctor/rules/compression_test.py | 9 +++-- 3 files changed, 64 insertions(+), 3 deletions(-) create mode 100644 .github/workflows/tests.yml create mode 100644 requirements.txt diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..343e8a3 --- /dev/null +++ b/.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/checkout@v4.0.0 + - name: Set up Python ${{ matrix.python-version }} + id: python + uses: actions/setup-python@v4.7.0 + 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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8f82d3f --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +psycopg2>=2.9.2 diff --git a/src/doctor/rules/compression_test.py b/src/doctor/rules/compression_test.py index be28b43..b74d4a7 100644 --- a/src/doctor/rules/compression_test.py +++ b/src/doctor/rules/compression_test.py @@ -34,11 +34,14 @@ class TestCompressionRules(unittest.TestCase): def setUp(self): """Set up unit tests for compression rules.""" - user = (os.getenv("PGUSER") or os.getlogin()) + user = os.getenv("PGUSER") 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, + 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",