pgmon/src/test_pgmon.py

import unittest

from datetime import datetime, timedelta
import tempfile

import logging

import pgmon

# Silence most logging output
logging.disable(logging.CRITICAL)


class TestPgmonMethods(unittest.TestCase):
    ##
    # update_deep
    ##
    def test_update_deep__empty_cases(self):
        # Test empty dict cases
        d1 = {}
        d2 = {}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {})
        self.assertEqual(d2, {})

        d1 = {"a": 1}
        d2 = {}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"a": 1})
        self.assertEqual(d2, {})

        d1 = {}
        d2 = {"a": 1}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"a": 1})
        self.assertEqual(d2, d1)

    def test_update_deep__scalars(self):
        # Test adding/updating scalar values
        d1 = {"foo": 1, "bar": "text", "hello": "world"}
        d2 = {"foo": 2, "baz": "blah"}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"foo": 2, "bar": "text", "baz": "blah", "hello": "world"})
        self.assertEqual(d2, {"foo": 2, "baz": "blah"})

    def test_update_deep__lists(self):
        # Test adding to lists
        d1 = {"lst1": []}
        d2 = {"lst1": [1, 2]}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"lst1": [1, 2]})
        self.assertEqual(d2, d1)

        d1 = {"lst1": [1, 2]}
        d2 = {"lst1": []}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"lst1": [1, 2]})
        self.assertEqual(d2, {"lst1": []})

        d1 = {"lst1": [1, 2, 3]}
        d2 = {"lst1": [3, 4]}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"lst1": [1, 2, 3, 3, 4]})
        self.assertEqual(d2, {"lst1": [3, 4]})

        # Lists of objects
        d1 = {"lst1": [{"id": 1}, {"id": 2}, {"id": 3}]}
        d2 = {"lst1": [{"id": 3}, {"id": 4}]}
        pgmon.update_deep(d1, d2)
        self.assertEqual(
            d1, {"lst1": [{"id": 1}, {"id": 2}, {"id": 3}, {"id": 3}, {"id": 4}]}
        )
        self.assertEqual(d2, {"lst1": [{"id": 3}, {"id": 4}]})

        # Nested lists
        d1 = {"obj1": {"l1": [1, 2]}}
        d2 = {"obj1": {"l1": [3, 4]}}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"obj1": {"l1": [1, 2, 3, 4]}})
        self.assertEqual(d2, {"obj1": {"l1": [3, 4]}})

    def test_update_deep__dicts(self):
        # Test adding to lists
        d1 = {"obj1": {}}
        d2 = {"obj1": {"a": 1, "b": 2}}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"obj1": {"a": 1, "b": 2}})
        self.assertEqual(d2, d1)

        d1 = {"obj1": {"a": 1, "b": 2}}
        d2 = {"obj1": {}}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"obj1": {"a": 1, "b": 2}})
        self.assertEqual(d2, {"obj1": {}})

        d1 = {"obj1": {"a": 1, "b": 2}}
        d2 = {"obj1": {"a": 5, "c": 12}}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"obj1": {"a": 5, "b": 2, "c": 12}})
        self.assertEqual(d2, {"obj1": {"a": 5, "c": 12}})

        # Nested dicts
        d1 = {"obj1": {"d1": {"a": 1, "b": 2}}}
        d2 = {"obj1": {"d1": {"a": 5, "c": 12}}}
        pgmon.update_deep(d1, d2)
        self.assertEqual(d1, {"obj1": {"d1": {"a": 5, "b": 2, "c": 12}}})
        self.assertEqual(d2, {"obj1": {"d1": {"a": 5, "c": 12}}})

    def test_update_deep__types(self):
        # Test mismatched types
        d1 = {"foo": 5}
        d2 = None
        self.assertRaises(TypeError, pgmon.update_deep, d1, d2)

        d1 = None
        d2 = {"foo": 5}
        self.assertRaises(TypeError, pgmon.update_deep, d1, d2)

        # Nested mismatched types
        d1 = {"foo": [1, 2]}
        d2 = {"foo": {"a": 7}}
        self.assertRaises(TypeError, pgmon.update_deep, d1, d2)

    ##
    # get_pool
    ##

    def test_get_pool__simple(self):
        # Just get a pool in a normal case
        pgmon.config.update(pgmon.default_config)
        pool = pgmon.get_pool("postgres")
        self.assertIsNotNone(pool)

    def test_get_pool__unhappy(self):
        # Test getting an unhappy database pool
        pgmon.config.update(pgmon.default_config)
        pgmon.unhappy_cooldown["postgres"] = datetime.now() + timedelta(60)
        self.assertRaises(pgmon.UnhappyDBError, pgmon.get_pool, "postgres")

        # Test getting a different database when there's an unhappy one
        pool = pgmon.get_pool("template0")
        self.assertIsNotNone(pool)

    ##
    # handle_connect_failure
    ##

    def test_handle_connect_failure__simple(self):
        # Test adding to an empty unhappy list
        pgmon.config.update(pgmon.default_config)
        pgmon.unhappy_cooldown = {}
        pool = pgmon.get_pool("postgres")
        pgmon.handle_connect_failure(pool)
        self.assertGreater(pgmon.unhappy_cooldown["postgres"], datetime.now())

        # Test adding another database
        pool = pgmon.get_pool("template0")
        pgmon.handle_connect_failure(pool)
        self.assertGreater(pgmon.unhappy_cooldown["postgres"], datetime.now())
        self.assertGreater(pgmon.unhappy_cooldown["template0"], datetime.now())
        self.assertEqual(len(pgmon.unhappy_cooldown), 2)

    ##
    # get_query
    ##

    def test_get_query__basic(self):
        # Test getting a query with one version
        metric = {"type": "value", "query": {0: "DEFAULT"}}
        self.assertEqual(pgmon.get_query(metric, 100000), "DEFAULT")

    def test_get_query__versions(self):
        metric = {"type": "value", "query": {0: "DEFAULT", 110000: "NEW"}}

        # Test getting the default version of a query with no lower bound and a newer version
        self.assertEqual(pgmon.get_query(metric, 100000), "DEFAULT")

        # Test getting the newer version of a query with no lower bound and a newer version for the newer version
        self.assertEqual(pgmon.get_query(metric, 110000), "NEW")

        # Test getting the newer version of a query with no lower bound and a newer version for an even newer version
        self.assertEqual(pgmon.get_query(metric, 160000), "NEW")

        # Test getting a version in bwtween two other versions
        metric = {"type": "value", "query": {0: "DEFAULT", 96000: "OLD", 110000: "NEW"}}
        self.assertEqual(pgmon.get_query(metric, 100000), "OLD")

    def test_get_query__missing_version(self):
        metric = {"type": "value", "query": {96000: "OLD", 110000: "NEW", 150000: ""}}

        # Test getting a metric that only exists for newer versions
        self.assertRaises(pgmon.MetricVersionError, pgmon.get_query, metric, 80000)

        # Test getting a metric that only exists for older versions
        self.assertRaises(pgmon.MetricVersionError, pgmon.get_query, metric, 160000)

    ##
    # read_config
    ##

    def test_read_config__simple(self):
        pgmon.config = {}

        # Test reading just a metric and using the defaults for everything else
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
# This is a comment!
metrics:
  test1:
    type: value
    query:
      0: TEST1
"""
                )

            pgmon.read_config(f"{tmpdirname}/config.yml")

        self.assertEqual(
            pgmon.config["max_pool_size"], pgmon.default_config["max_pool_size"]
        )
        self.assertEqual(pgmon.config["dbuser"], pgmon.default_config["dbuser"])

        pgmon.config = {}

        # Test reading a basic config
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
# This is a comment!
min_pool_size: 1
max_pool_size: 2
max_idle_time: 10
log_level: debug
dbuser: someone
dbhost: localhost
dbport: 5555
dbname: template0
pool_slot_timeout: 1
connect_timeout: 1
reconnect_cooldown: 15
version_check_period: 3600
metrics:
  test1:
    type: value
    query:
      0: TEST1
  test2:
    type: set
    query:
      0: TEST2
  test3:
    type: row
    query:
      0: TEST3
  test4:
    type: column
    query:
      0: TEST4
"""
                )

            pgmon.read_config(f"{tmpdirname}/config.yml")

        self.assertEqual(pgmon.config["dbuser"], "someone")
        self.assertEqual(pgmon.config["metrics"]["test1"]["type"], "value")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")
        self.assertEqual(pgmon.config["metrics"]["test2"]["query"][0], "TEST2")

    def test_read_config__include(self):
        pgmon.config = {}

        # Test reading a config that includes other files (absolute and relative paths, multiple levels)
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    f"""---
# This is a comment!
min_pool_size: 1
max_pool_size: 2
max_idle_time: 10
log_level: debug
pool_slot_timeout: 1
connect_timeout: 1
reconnect_cooldown: 15
version_check_period: 3600
include:
  - dbsettings.yml
  - {tmpdirname}/metrics.yml
"""
                )

            with open(f"{tmpdirname}/dbsettings.yml", "w") as f:
                f.write(
                    f"""---
dbuser: someone
dbhost: localhost
dbport: 5555
dbname: template0
"""
                )

            with open(f"{tmpdirname}/metrics.yml", "w") as f:
                f.write(
                    f"""---
metrics:
  test1:
    type: value
    query:
      0: TEST1
  test2:
    type: value
    query:
      0: TEST2
include:
  - more_metrics.yml
"""
                )

            with open(f"{tmpdirname}/more_metrics.yml", "w") as f:
                f.write(
                    f"""---
metrics:
  test3:
    type: value
    query:
      0: TEST3
"""
                )
            pgmon.read_config(f"{tmpdirname}/config.yml")

        self.assertEqual(pgmon.config["max_idle_time"], 10)
        self.assertEqual(pgmon.config["dbuser"], "someone")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")
        self.assertEqual(pgmon.config["metrics"]["test2"]["query"][0], "TEST2")
        self.assertEqual(pgmon.config["metrics"]["test3"]["query"][0], "TEST3")

    def test_read_config__reload(self):
        pgmon.config = {}

        # Test rereading a config to update an existing config
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
# This is a comment!
min_pool_size: 1
max_pool_size: 2
max_idle_time: 10
log_level: debug
dbuser: someone
dbhost: localhost
dbport: 5555
dbname: template0
pool_slot_timeout: 1
connect_timeout: 1
reconnect_cooldown: 15
version_check_period: 3600
metrics:
  test1:
    type: value
    query:
      0: TEST1
  test2:
    type: value
    query:
      0: TEST2
"""
                )

            pgmon.read_config(f"{tmpdirname}/config.yml")

            # Just make sure the first config was read
            self.assertEqual(len(pgmon.config["metrics"]), 2)

            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
# This is a comment!
min_pool_size: 7
metrics:
  test1:
    type: value
    query:
      0: NEW1
"""
                )

            pgmon.read_config(f"{tmpdirname}/config.yml")

        self.assertEqual(pgmon.config["min_pool_size"], 7)
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "NEW1")
        self.assertEqual(len(pgmon.config["metrics"]), 1)

    def test_read_config__query_file(self):
        pgmon.config = {}

        # Read a config file that reads a query from a file
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
metrics:
  test1:
    type: value
    query:
      0: file:some_query.sql
"""
                )

            with open(f"{tmpdirname}/some_query.sql", "w") as f:
                f.write("This is a query")

            pgmon.read_config(f"{tmpdirname}/config.yml")

        self.assertEqual(
            pgmon.config["metrics"]["test1"]["query"][0], "This is a query"
        )

    def test_read_config__invalid(self):
        pgmon.config = {}

        # For all of these tests, we start with a valid config and also ensure that
        #  it is not modified when a new config read fails
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
metrics:
  test1:
    type: value
    query:
      0: TEST1
"""
                )

            pgmon.read_config(f"{tmpdirname}/config.yml")

        # Just make sure the config was read
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test reading a nonexistant config file
        with tempfile.TemporaryDirectory() as tmpdirname:
            self.assertRaises(
                FileNotFoundError, pgmon.read_config, f"{tmpdirname}/missing.yml"
            )

        # Test reading an invalid config file
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """[default]
This looks a lot like an ini file to me

Or maybe a TOML?
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )

        # Test reading a config that includes an invalid file
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    type: value
    query:
      0: EVIL1
include:
  - missing_file.yml
"""
                )
            self.assertRaises(
                FileNotFoundError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test invalid log level
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
log_level: noisy
dbuser: evil
metrics:
  test1:
    type: value
    query:
      0: EVIL1
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test invalid query return type
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    type: lots_of_data
    query:
      0: EVIL1
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test invalid query dict type
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    type: lots_of_data
    query: EVIL1
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test incomplete metric: missing type
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    query:
      0: EVIL1
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test incomplete metric: missing queries
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    type: value
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test incomplete metric: empty queries
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    type: value
    query: {}
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test incomplete metric: query dict is None
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    type: value
    query:
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test reading a config with no metrics
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test reading a query defined in a file but the file is missing
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    type: value
    query: 
      0: file:missing.sql
"""
                )
            self.assertRaises(
                FileNotFoundError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")

        # Test invalid query versions
        with tempfile.TemporaryDirectory() as tmpdirname:
            with open(f"{tmpdirname}/config.yml", "w") as f:
                f.write(
                    """---
dbuser: evil
metrics:
  test1:
    type: value
    query: 
      default: EVIL1
"""
                )
            self.assertRaises(
                pgmon.ConfigError, pgmon.read_config, f"{tmpdirname}/config.yml"
            )
        self.assertEqual(pgmon.config["dbuser"], "postgres")
        self.assertEqual(pgmon.config["metrics"]["test1"]["query"][0], "TEST1")