Added tests to verify json dump customisation

diff --git a/tests/types/test_json.py b/tests/types/test_json.py
index 8cb401fc3db13aec66f31ff028d9dc2682c699ae..9c003cd58b78bb1cc0e01afd944866a394cef740 100644 (file)
--- a/tests/types/test_json.py
+++ b/tests/types/test_json.py
@@ -2,6 +2,7 @@ import json
 
 import pytest
 
+import psycopg3.types.json
 from psycopg3.types.json import Json, JsonB
 from psycopg3.adapt import Format
 
@@ -37,3 +38,37 @@ def test_jsonb_dump(conn, val, fmt_in):
     cur = conn.cursor()
     cur.execute(f"select {ph} = %s::jsonb", (JsonB(obj), val))
     assert cur.fetchone()[0] is True
+
+
+def my_dumps(obj):
+    obj["baz"] = "qux"
+    return json.dumps(obj)
+
+
+@pytest.mark.parametrize("fmt_in", [Format.TEXT, Format.BINARY])
+@pytest.mark.parametrize("wrapper", ["Json", "JsonB"])
+def test_json_dump_customise(conn, wrapper, fmt_in):
+    ph = "%s" if fmt_in == Format.TEXT else "%b"
+    wrapper = getattr(psycopg3.types.json, wrapper)
+    obj = {"foo": "bar"}
+    cur = conn.cursor()
+    cur.execute(
+        f"select {ph}->>'baz' = 'qux'", (wrapper(obj, dumps=my_dumps),)
+    )
+    assert cur.fetchone()[0] is True
+
+
+@pytest.mark.parametrize("fmt_in", [Format.TEXT, Format.BINARY])
+@pytest.mark.parametrize("wrapper", ["Json", "JsonB"])
+def test_json_dump_subclass(conn, wrapper, fmt_in):
+    ph = "%s" if fmt_in == Format.TEXT else "%b"
+    wrapper = getattr(psycopg3.types.json, wrapper)
+
+    class MyWrapper(wrapper):
+        def dumps(self):
+            return my_dumps(self.obj)
+
+    obj = {"foo": "bar"}
+    cur = conn.cursor()
+    cur.execute(f"select {ph}->>'baz' = 'qux'", (MyWrapper(obj),))
+    assert cur.fetchone()[0] is True