From f13d04d1695ba01109ea1a32b6a3fbad75ebfd30 Mon Sep 17 00:00:00 2001
From: Robin Howard <robin@cipherstash.com>
Date: Wed, 9 Jul 2025 15:55:12 +1000
Subject: [PATCH] test(python): Adds contained-by ($1 <@ a_jsonb_column) tests
 for the Python integration test suite.

---
 tests/python/tests/test_jsonb_ops.py | 88 ++++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)
 create mode 100644 tests/python/tests/test_jsonb_ops.py

diff --git a/tests/python/tests/test_jsonb_ops.py b/tests/python/tests/test_jsonb_ops.py
new file mode 100644
index 00000000..823180e8
--- /dev/null
+++ b/tests/python/tests/test_jsonb_ops.py
@@ -0,0 +1,88 @@
+import json
+import os
+import psycopg
+import random
+
+conn_params = {
+    "user": os.environ.get("CS_DATABASE__USERNAME"),
+    "password": os.environ.get("CS_DATABASE__PASSWORD"),
+    "dbname": os.environ.get("CS_DATABASE__NAME"),
+    "host": os.environ.get("CS_DATABASE__HOST"),
+    "port": 6432,
+}
+
+connection_str = psycopg.conninfo.make_conninfo(**conn_params)
+
+print("Connection to Proxy with {}".format(connection_str))
+
+
+def test_jsonb_contained_by():
+    val = {"key": "value"}
+    column = "encrypted_jsonb"
+    select_fragment = "%s <@ encrypted_jsonb"
+    tests = [
+        (val, True),
+        ({"key": "different value"}, False)
+    ]
+
+    for (param, expected) in tests:
+        param = json.dumps(param)
+
+        execute(json.dumps(val), column,
+                select_fragment=select_fragment,
+                select_params=[param],
+                expected=expected)
+
+        execute(json.dumps(val).encode(), column,
+                select_fragment=select_fragment,
+                select_params=[param.encode()],
+                expected=expected,
+                binary=True)
+
+        execute(json.dumps(val).encode(), column,
+                select_fragment=select_fragment,
+                select_params=[param.encode()],
+                expected=expected,
+                binary=True, prepare=True)
+
+        execute(json.dumps(val), column,
+                select_fragment=select_fragment,
+                select_params=[param],
+                expected=expected,
+                binary=False, prepare=True)
+
+
+def make_id():
+    return random.randrange(1, 1000000000)
+
+
+def execute(val, column, binary=None, prepare=None, expected=None,
+            select_fragment=None, select_params=[]):
+    with psycopg.connect(connection_str, autocommit=True) as conn:
+
+        with conn.cursor() as cursor:
+
+            with conn.transaction():
+                id = make_id()
+
+                print("Testing {} Binary: {} Prepare: {}".format(
+                    column, binary, prepare))
+
+                sql = "INSERT INTO encrypted (id, {}) VALUES (%s, %s)".format(
+                    column)
+
+                cursor.execute(sql, [id, val], binary=binary, prepare=prepare)
+
+                sql = "SELECT id, {} FROM encrypted WHERE id = %s".format(
+                    select_fragment)
+                cursor.execute(
+                    sql, (select_params + [id]),
+                    binary=binary, prepare=prepare)
+
+                row = cursor.fetchone()
+
+                (result_id, result) = row
+                expected_result = expected if expected is not None else val
+
+                assert result_id == id
+                assert result == expected_result