Interpret incoming JSON as a string instead of via json.loads

Closes #163

1 files changed, 21 insertions, 2 deletions

diff --git a/tests/utils.py b/tests/utils.py
index 9b7376f1..b755f78e 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,6 +1,8 @@
 import pytest
 import psycopg2
+import psycopg2.extras
 from pgcli.main import format_output
+from pgcli.pgexecute import register_json_typecasters
 
 # TODO: should this be somehow be divined from environment?
 POSTGRES_USER, POSTGRES_HOST = 'postgres', 'localhost'
@@ -11,17 +13,34 @@ def db_connection(dbname=None):
     conn.autocommit = True
     return conn
 
+
 try:
-    db_connection()
+    conn = db_connection()
     CAN_CONNECT_TO_DB = True
+    SERVER_VERSION = conn.server_version
+    json_types = register_json_typecasters(conn, lambda x: x)
+    JSON_AVAILABLE = 'json' in json_types
+    JSONB_AVAILABLE = 'jsonb' in json_types
 except:
-    CAN_CONNECT_TO_DB = False
+    CAN_CONNECT_TO_DB = JSON_AVAILABLE = JSONB_AVAILABLE = False
+    SERVER_VERSION = 0
+
 
 dbtest = pytest.mark.skipif(
     not CAN_CONNECT_TO_DB,
     reason="Need a postgres instance at localhost accessible by user 'postgres'")
 
 
+requires_json = pytest.mark.skipif(
+    not JSON_AVAILABLE,
+    reason='Postgres server unavailable or json type not defined')
+
+
+requires_jsonb = pytest.mark.skipif(
+    not JSONB_AVAILABLE,
+    reason='Postgres server unavailable or jsonb type not defined')
+
+
 def create_db(dbname):
     with db_connection().cursor() as cur:
         try: