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