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')