1 files changed, 108 insertions, 0 deletions

diff --git a/tests/test_dbs.py b/tests/test_dbs.py
new file mode 100644
index 0000000..157c423
--- /dev/null
+++ b/tests/test_dbs.py
@@ -0,0 +1,108 @@
+from .helper import *
+import dbs, main
+
+@unittest.skipUnless("unit" in limits, "Unit skipped")
+class DbsTest(WebMockTestCase):
+    @mock.patch.object(dbs, "psycopg2")
+    def test_connect_psql(self, psycopg2):
+        args = main.configargparse.Namespace(**{
+            "db_host": "host",
+            "db_port": "port",
+            "db_user": "user",
+            "db_password": "password",
+            "db_database": "database",
+            })
+        psycopg2.connect.return_value = "pg_connection"
+        dbs.connect_psql(args)
+
+        psycopg2.connect.assert_called_once_with(host="host",
+                port="port", user="user", password="password",
+                database="database")
+        self.assertEqual("pg_connection", dbs.psql)
+        with self.assertRaises(AttributeError):
+            args.db_password
+
+        psycopg2.connect.reset_mock()
+        args = main.configargparse.Namespace(**{
+            "db_host": "host",
+            "db_port": "port",
+            "db_user": "user",
+            "db_password": "password",
+            "db_database": "database",
+            })
+        dbs.connect_psql(args)
+        psycopg2.connect.assert_not_called()
+
+    @mock.patch.object(dbs, "_redis")
+    def test_connect_redis(self, redis):
+        with self.subTest(redis_host="tcp"):
+            args = main.configargparse.Namespace(**{
+                "redis_host": "host",
+                "redis_port": "port",
+                "redis_database": "database",
+                })
+            redis.Redis.return_value = "redis_connection"
+            dbs.connect_redis(args)
+
+            redis.Redis.assert_called_once_with(host="host",
+                    port="port", db="database")
+            self.assertEqual("redis_connection", dbs.redis)
+            with self.assertRaises(AttributeError):
+                args.redis_database
+
+            redis.Redis.reset_mock()
+            args = main.configargparse.Namespace(**{
+                "redis_host": "host",
+                "redis_port": "port",
+                "redis_database": "database",
+                })
+            dbs.connect_redis(args)
+            redis.Redis.assert_not_called()
+
+        dbs.redis = None
+        with self.subTest(redis_host="socket"):
+            args = main.configargparse.Namespace(**{
+                "redis_host": "/run/foo",
+                "redis_port": "port",
+                "redis_database": "database",
+                })
+            redis.Redis.return_value = "redis_socket"
+            dbs.connect_redis(args)
+
+            redis.Redis.assert_called_once_with(unix_socket_path="/run/foo", db="database")
+            self.assertEqual("redis_socket", dbs.redis)
+
+    def test_redis_connected(self):
+        with self.subTest(redis=None):
+            dbs.redis = None
+            self.assertFalse(dbs.redis_connected())
+
+        with self.subTest(redis="mocked_true"):
+            dbs.redis = mock.Mock()
+            dbs.redis.ping.return_value = True
+            self.assertTrue(dbs.redis_connected())
+
+        with self.subTest(redis="mocked_false"):
+            dbs.redis = mock.Mock()
+            dbs.redis.ping.return_value = False
+            self.assertFalse(dbs.redis_connected())
+
+        with self.subTest(redis="mocked_raise"):
+            dbs.redis = mock.Mock()
+            dbs.redis.ping.side_effect = Exception("bouh")
+            self.assertFalse(dbs.redis_connected())
+
+    def test_psql_connected(self):
+        with self.subTest(psql=None):
+            dbs.psql = None
+            self.assertFalse(dbs.psql_connected())
+
+        with self.subTest(psql="connected"):
+            dbs.psql = mock.Mock()
+            dbs.psql.closed = 0
+            self.assertTrue(dbs.psql_connected())
+
+        with self.subTest(psql="not connected"):
+            dbs.psql = mock.Mock()
+            dbs.psql.closed = 3
+            self.assertFalse(dbs.psql_connected())