Test: Improve persistence test handling

Sqlite uses contextlib for opening connections
Checks can optionally pass an existing connection so the db only needs
to be opened once even in the full check_db case
This commit is contained in:
Roger A. Light 2026-04-30 19:12:54 +01:00 committed by Roger Light
parent 06b22f9bcf
commit 1cd83de189
2 changed files with 210 additions and 235 deletions

View file

@ -110,76 +110,82 @@ def check_db(
qos: int = 1, qos: int = 1,
check_session_expiry_time: bool = True, check_session_expiry_time: bool = True,
): ):
count_list = [v for v in client_msg_counts.values() if v is not None] + [0]
num_base_msgs = max(count_list)
num_subscriptions = sum(1 for c in client_msg_counts.values() if c is not None)
num_client_msgs_out = sum(count_list)
persist_help.check_counts(
port,
clients=len(client_msg_counts),
client_msgs_out=num_client_msgs_out,
base_msgs=num_base_msgs if num_base_msgs > 0 or retain_end == 0 else 1,
retain_msgs=1 if retain_end > 0 else 0,
subscriptions=num_subscriptions,
)
# Check client with persist_help.get_connection(port) as con:
for client_id, num_messages_for_client in client_msg_counts.items(): count_list = [v for v in client_msg_counts.values() if v is not None] + [0]
persist_help.check_client( num_base_msgs = max(count_list)
num_subscriptions = sum(1 for c in client_msg_counts.values() if c is not None)
num_client_msgs_out = sum(count_list)
persist_help.check_counts(
port, port,
client_id, clients=len(client_msg_counts),
username=username, client_msgs_out=num_client_msgs_out,
will_delay_time=0, base_msgs=num_base_msgs if num_base_msgs > 0 or retain_end == 0 else 1,
session_expiry_time=60 if check_session_expiry_time else None, retain_msgs=1 if retain_end > 0 else 0,
listener_port=None, # persist-lmdb reset listener port to 0 on disconnect subscriptions=num_subscriptions,
max_packet_size=0, connection=con,
max_qos=2,
retain_available=1,
session_expiry_interval=60,
will_delay_interval=0,
) )
# Check subscription
if num_messages_for_client is not None:
persist_help.check_subscription(port, client_id, subscription_topic, qos, 0)
# Check stored message # Check client
for i in range(num_base_msgs):
msg_id = num_published_msgs - num_base_msgs + i
payload = f"queued message {msg_id:3}"
payload_b = payload.encode("UTF-8")
mid = 10 + msg_id
store_id = persist_help.check_base_msg(
port,
message_expiry,
subscription_topic,
payload_b,
publisher_id,
username,
len(payload_b),
mid,
port,
qos,
retain=1 if i < retain_end else 0,
idx=i,
)
# Check client msg
for client_id, num_messages_for_client in client_msg_counts.items(): for client_id, num_messages_for_client in client_msg_counts.items():
if num_messages_for_client is None: persist_help.check_client(
continue
client_msg_start = num_published_msgs - num_messages_for_client
if msg_id < client_msg_start:
continue
cmsg_id = 1 + msg_id - client_msg_start
subscriber_mid = cmsg_id
persist_help.check_client_msg(
port, port,
client_id, client_id,
cmsg_id, username=username,
store_id, will_delay_time=0,
0, session_expiry_time=60 if check_session_expiry_time else None,
persist_help.dir_out, listener_port=None, # persist-lmdb reset listener port to 0 on disconnect
subscriber_mid, max_packet_size=0,
qos, max_qos=2,
0, retain_available=1,
persist_help.ms_queued, session_expiry_interval=60,
will_delay_interval=0,
connection=con,
) )
# Check subscription
if num_messages_for_client is not None:
persist_help.check_subscription(port, client_id, subscription_topic, qos, 0, connection=con)
# Check stored message
for i in range(num_base_msgs):
msg_id = num_published_msgs - num_base_msgs + i
payload = f"queued message {msg_id:3}"
payload_b = payload.encode("UTF-8")
mid = 10 + msg_id
store_id = persist_help.check_base_msg(
port,
message_expiry,
subscription_topic,
payload_b,
publisher_id,
username,
len(payload_b),
mid,
port,
qos,
retain=1 if i < retain_end else 0,
idx=i,
connection=con,
)
# Check client msg
for client_id, num_messages_for_client in client_msg_counts.items():
if num_messages_for_client is None:
continue
client_msg_start = num_published_msgs - num_messages_for_client
if msg_id < client_msg_start:
continue
cmsg_id = 1 + msg_id - client_msg_start
subscriber_mid = cmsg_id
persist_help.check_client_msg(
port,
client_id,
cmsg_id,
store_id,
0,
persist_help.dir_out,
subscriber_mid,
qos,
0,
persist_help.ms_queued,
connection=con,
)

View file

@ -1,4 +1,5 @@
import os import os
from contextlib import contextmanager
from pathlib import Path from pathlib import Path
import sqlite3 import sqlite3
import mosq_paths import mosq_paths
@ -23,6 +24,17 @@ ms_wait_for_pubcomp = 9
ms_send_pubrec = 10 ms_send_pubrec = 10
ms_queued = 11 ms_queued = 11
@contextmanager
def get_connection(port, connection=None):
if connection is not None:
yield connection
else:
con = sqlite3.connect(Path(str(port), 'mosquitto.sqlite3'))
try:
yield con
finally:
con.close()
def write_config(filename, port, additional_config_entries: dict = {}): def write_config(filename, port, additional_config_entries: dict = {}):
with open(filename, "w") as f: with open(filename, "w") as f:
@ -119,21 +131,18 @@ def cleanup(port):
@retry() @retry()
def check_version_infos(port, database_schema_version): def check_version_infos(port, database_schema_version):
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port) as con:
cur = con.cursor() row = con.execute(
cur.execute( "SELECT major,minor,patch FROM version_info WHERE component = 'database_schema';"
"SELECT major,minor,patch FROM version_info WHERE component = 'database_schema';" ).fetchone()
)
row = cur.fetchone()
if len(row) != len(database_schema_version): if len(row) != len(database_schema_version):
raise ValueError("Could not fetch db version info from DB") raise ValueError("Could not fetch db version info from DB")
for i in range(len(row)): for i in range(len(row)):
if row[i] != database_schema_version[i]: if row[i] != database_schema_version[i]:
raise ValueError( raise ValueError(
f"DB version info {'.'.join([str(v) for v in row])} != expected {'.'.join([str(v) for v in database_schema_version])}" f"DB version info {'.'.join([str(v) for v in row])} != expected {'.'.join([str(v) for v in database_schema_version])}"
) )
con.close()
@retry() @retry()
@ -145,54 +154,44 @@ def check_counts(
base_msgs=0, base_msgs=0,
retain_msgs=0, retain_msgs=0,
subscriptions=0, subscriptions=0,
wills=None wills=None,
connection = None,
): ):
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port, connection) as con:
cur = con.cursor() row = con.execute("SELECT COUNT(*) FROM clients").fetchone()
cur.execute("SELECT COUNT(*) FROM clients") if row[0] != clients:
row = cur.fetchone() raise ValueError("Found %d clients, expected %d" % (row[0], clients))
if row[0] != clients:
raise ValueError("Found %d clients, expected %d" % (row[0], clients))
cur.execute("SELECT COUNT(*) FROM client_msgs WHERE direction=0") row = con.execute("SELECT COUNT(*) FROM client_msgs WHERE direction=0").fetchone()
row = cur.fetchone() if row[0] != client_msgs_in:
if row[0] != client_msgs_in: raise ValueError(
raise ValueError( "Found %d client_msgs_in, expected %d" % (row[0], client_msgs_in)
"Found %d client_msgs_in, expected %d" % (row[0], client_msgs_in) )
)
cur.execute("SELECT COUNT(*) FROM client_msgs WHERE direction=1") row = con.execute("SELECT COUNT(*) FROM client_msgs WHERE direction=1").fetchone()
row = cur.fetchone() if row[0] != client_msgs_out:
if row[0] != client_msgs_out: raise ValueError(
raise ValueError( "Found %d client_msgs_out, expected %d" % (row[0], client_msgs_out)
"Found %d client_msgs_out, expected %d" % (row[0], client_msgs_out) )
)
cur.execute("SELECT COUNT(*) FROM subscriptions") row = con.execute("SELECT COUNT(*) FROM subscriptions").fetchone()
row = cur.fetchone() if row[0] != subscriptions:
if row[0] != subscriptions: raise ValueError(
raise ValueError( "Found %d subscriptions, expected %d" % (row[0], subscriptions)
"Found %d subscriptions, expected %d" % (row[0], subscriptions) )
)
cur.execute("SELECT COUNT(*) FROM base_msgs") row = con.execute("SELECT COUNT(*) FROM base_msgs").fetchone()
row = cur.fetchone() if row[0] != base_msgs:
if row[0] != base_msgs: raise ValueError("Found %d base_msgs, expected %d" % (row[0], base_msgs))
raise ValueError("Found %d base_msgs, expected %d" % (row[0], base_msgs))
cur.execute("SELECT COUNT(*) FROM retains") row = con.execute("SELECT COUNT(*) FROM retains").fetchone()
row = cur.fetchone() if row[0] != retain_msgs:
if row[0] != retain_msgs: raise ValueError("Found %d retain_msgs, expected %d" % (row[0], retain_msgs))
raise ValueError("Found %d retain_msgs, expected %d" % (row[0], retain_msgs))
if wills is not None: if wills is not None:
cur.execute("SELECT COUNT(*) FROM wills") row = con.execute("SELECT COUNT(*) FROM wills").fetchone()
row = cur.fetchone() if row[0] != wills:
if row[0] != wills: raise ValueError("Found %d wills, expected %d" % (row[0], wills))
raise ValueError("Found %d wills, expected %d" % (row[0], wills))
con.close()
@retry() @retry()
@ -208,76 +207,72 @@ def check_client(
retain_available, retain_available,
session_expiry_interval, session_expiry_interval,
will_delay_interval, will_delay_interval,
connection=None,
): ):
# "Fix" the infinite session expiry interval as mangled by an int32 conversion. # "Fix" the infinite session expiry interval as mangled by an int32 conversion.
if session_expiry_interval == 4294967295: if session_expiry_interval == 4294967295:
session_expiry_interval = -1 session_expiry_interval = -1
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port, connection) as con:
cur = con.cursor() row = con.execute(
cur.execute( "SELECT client_id, username, will_delay_time, session_expiry_time, "
"SELECT client_id, username, will_delay_time, session_expiry_time, " + "listener_port, max_packet_size, max_qos, retain_available, "
+ "listener_port, max_packet_size, max_qos, retain_available, " + "session_expiry_interval, will_delay_interval "
+ "session_expiry_interval, will_delay_interval " + "FROM clients "
+ "FROM clients " + f"WHERE client_id = '{client_id}'"
+ f"WHERE client_id = '{client_id}'" ).fetchone()
)
row = cur.fetchone()
if row is None: if row is None:
raise ValueError(f"Cannot find client {client_id} in db") raise ValueError(f"Cannot find client {client_id} in db")
if row[0] != client_id: if row[0] != client_id:
raise ValueError("Invalid client_id %s / %s" % (row[0], client_id)) raise ValueError("Invalid client_id %s / %s" % (row[0], client_id))
if username is not None and row[1] != username: if username is not None and row[1] != username:
raise ValueError("Invalid username %s / %s" % (row[1], username)) raise ValueError("Invalid username %s / %s" % (row[1], username))
if (will_delay_time == 0 and row[2] != 0) or (will_delay_time != 0 and row[2] == 0): if (will_delay_time == 0 and row[2] != 0) or (will_delay_time != 0 and row[2] == 0):
raise ValueError("Invalid will_delay_time %d / %d" % (row[2], will_delay_time)) raise ValueError("Invalid will_delay_time %d / %d" % (row[2], will_delay_time))
if session_expiry_time and ( if session_expiry_time and (
(session_expiry_time == 0 and row[3] != 0) (session_expiry_time == 0 and row[3] != 0)
or (session_expiry_time != 0 and row[3] == 0) or (session_expiry_time != 0 and row[3] == 0)
): ):
raise ValueError( raise ValueError(
"Invalid session_expiry_time %d / %d for client %s" "Invalid session_expiry_time %d / %d for client %s"
% (row[3], session_expiry_time, client_id) % (row[3], session_expiry_time, client_id)
) )
if listener_port is not None and row[4] != listener_port: if listener_port is not None and row[4] != listener_port:
raise ValueError("Invalid listener_port %d / %d" % (row[4], listener_port)) raise ValueError("Invalid listener_port %d / %d" % (row[4], listener_port))
if row[5] != max_packet_size: if row[5] != max_packet_size:
raise ValueError("Invalid max_packet_size %d / %d" % (row[5], max_packet_size)) raise ValueError("Invalid max_packet_size %d / %d" % (row[5], max_packet_size))
if row[6] != max_qos: if row[6] != max_qos:
raise ValueError("Invalid max_qos %d / %d" % (row[6], max_qos)) raise ValueError("Invalid max_qos %d / %d" % (row[6], max_qos))
if row[7] != retain_available: if row[7] != retain_available:
raise ValueError( raise ValueError(
"Invalid retain_available %d / %d" % (row[7], retain_available) "Invalid retain_available %d / %d" % (row[7], retain_available)
) )
if row[8] != session_expiry_interval: if row[8] != session_expiry_interval:
raise ValueError( raise ValueError(
"Invalid session_expiry_interval %d / %d" "Invalid session_expiry_interval %d / %d"
% (row[8], session_expiry_interval) % (row[8], session_expiry_interval)
) )
if row[9] != will_delay_interval: if row[9] != will_delay_interval:
raise ValueError( raise ValueError(
"Invalid will_delay_interval %d / %d" % (row[9], will_delay_interval) "Invalid will_delay_interval %d / %d" % (row[9], will_delay_interval)
) )
con.close()
def modify_client(port: int, client_id: str, sub_expiry_time: int): def modify_client(port: int, client_id: str, sub_expiry_time: int, connection=None):
num_modified_rows = 0 num_modified_rows = 0
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port, connection) as con:
try: cur = con.execute(
cur = con.cursor()
cur.execute(
"UPDATE clients" "UPDATE clients"
+ f" SET session_expiry_time = session_expiry_time - {sub_expiry_time}" + f" SET session_expiry_time = session_expiry_time - {sub_expiry_time}"
+ f" WHERE client_id = ?", + f" WHERE client_id = ?",
@ -285,60 +280,52 @@ def modify_client(port: int, client_id: str, sub_expiry_time: int):
) )
num_modified_rows = cur.rowcount num_modified_rows = cur.rowcount
con.commit() con.commit()
finally:
con.close()
return num_modified_rows return num_modified_rows
@retry() @retry()
def check_subscription( def check_subscription(
port, client_id, topic, subscription_options, subscription_identifier port, client_id, topic, subscription_options, subscription_identifier, connection=None
): ):
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port, connection) as con:
cur = con.cursor() row = con.execute(
cur.execute( "SELECT client_id, topic, subscription_options, subscription_identifier "
"SELECT client_id, topic, subscription_options, subscription_identifier " + "FROM subscriptions "
+ "FROM subscriptions " + f"WHERE client_id = '{client_id}'"
+ f"WHERE client_id = '{client_id}'" ).fetchone()
)
row = cur.fetchone()
if row is None: if row is None:
raise ValueError(f"Cannot find client {client_id} in db") raise ValueError(f"Cannot find client {client_id} in db")
if row[0] != client_id: if row[0] != client_id:
raise ValueError("Invalid client_id %s / %s" % (row[0], client_id)) raise ValueError("Invalid client_id %s / %s" % (row[0], client_id))
if row[1] != topic: if row[1] != topic:
raise ValueError("Invalid topic %s / %s" % (row[1], topic)) raise ValueError("Invalid topic %s / %s" % (row[1], topic))
if row[2] != subscription_options: if row[2] != subscription_options:
raise ValueError( raise ValueError(
"Invalid subscription_options %d / %d" % (row[2], subscription_options) "Invalid subscription_options %d / %d" % (row[2], subscription_options)
) )
if row[3] != subscription_identifier: if row[3] != subscription_identifier:
raise ValueError( raise ValueError(
"Invalid subscription_identifier %d / %d" "Invalid subscription_identifier %d / %d"
% (row[3], subscription_identifier) % (row[3], subscription_identifier)
) )
con.close()
@retry() @retry()
def check_client_msg( def check_client_msg(
port, client_id, cmsg_id, store_id, dup, direction, mid, qos, retain, state port, client_id, cmsg_id, store_id, dup, direction, mid, qos, retain, state, connection=None
): ):
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port, connection) as con:
try: row = con.execute(
cur = con.cursor()
cur.execute(
"SELECT client_id,cmsg_id,store_id,dup,direction,mid,qos,retain,state " "SELECT client_id,cmsg_id,store_id,dup,direction,mid,qos,retain,state "
+ "FROM client_msgs " + "FROM client_msgs "
+ f"WHERE client_id = '{client_id}' AND cmsg_id = {cmsg_id}" + f"WHERE client_id = '{client_id}' AND cmsg_id = {cmsg_id}"
) ).fetchone()
row = cur.fetchone()
msg_id = f"client_id={client_id},cmsg_id={cmsg_id}" msg_id = f"client_id={client_id},cmsg_id={cmsg_id}"
if row is None: if row is None:
@ -390,8 +377,6 @@ def check_client_msg(
raise ValueError( raise ValueError(
"Invalid state %d / %d for message %s" % (row[8], state, msg_id) "Invalid state %d / %d for message %s" % (row[8], state, msg_id)
) )
finally:
con.close()
@retry() @retry()
@ -408,11 +393,10 @@ def check_base_msg(
qos, qos,
retain, retain,
idx=0, idx=0,
connection=None,
): ):
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port, connection) as con:
try: cur = con.execute(
cur = con.cursor()
cur.execute(
"SELECT store_id,expiry_time,topic,payload,source_id,source_username, " "SELECT store_id,expiry_time,topic,payload,source_id,source_username, "
+ "payloadlen, source_mid, source_port, qos, retain " + "payloadlen, source_mid, source_port, qos, retain "
+ "FROM base_msgs " + "FROM base_msgs "
@ -457,10 +441,6 @@ def check_base_msg(
if row[10] != retain: if row[10] != retain:
raise ValueError("Invalid retain %d / %d" % (row[10], retain)) raise ValueError("Invalid retain %d / %d" % (row[10], retain))
except ValueError as err:
raise ValueError(str(err) + f" at index {idx}") from err
finally:
con.close()
return row[0] return row[0]
@ -468,31 +448,25 @@ def check_base_msg(
def modify_base_msgs( def modify_base_msgs(
port: int, port: int,
sub_expiry_time: int, sub_expiry_time: int,
connection=None,
): ):
num_modified_rows = 0 num_modified_rows = 0
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port, connection) as con:
try: cur = con.execute(
cur = con.cursor()
cur.execute(
"UPDATE base_msgs" + f" SET expiry_time = expiry_time - {sub_expiry_time}" "UPDATE base_msgs" + f" SET expiry_time = expiry_time - {sub_expiry_time}"
) )
num_modified_rows = cur.rowcount num_modified_rows = cur.rowcount
con.commit() con.commit()
finally:
con.close()
return num_modified_rows return num_modified_rows
@retry() @retry()
def check_retain(port, topic, store_id): def check_retain(port, topic, store_id):
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) as con:
cur = con.cursor() row = con.execute("SELECT store_id FROM retains WHERE topic=?", (topic,)).fetchone()
cur.execute("SELECT store_id FROM retains WHERE topic=?", (topic,))
row = cur.fetchone()
if row[0] != store_id: if row[0] != store_id:
raise ValueError("Invalid store_id %d / %d" % (row[0], store_id)) raise ValueError("Invalid store_id %d / %d" % (row[0], store_id))
con.close()
@retry() @retry()
@ -505,9 +479,9 @@ def check_will(
retain: int, retain: int,
properties: str, properties: str,
idx=0, idx=0,
connection=None,
): ):
con = sqlite3.connect(Path(str(port), "mosquitto.sqlite3")) with get_connection(port, connection) as con:
try:
cur = con.cursor() cur = con.cursor()
cur.execute( cur.execute(
"SELECT client_id,topic,payload,payloadlen,qos,retain,properties " "SELECT client_id,topic,payload,payloadlen,qos,retain,properties "
@ -540,9 +514,4 @@ def check_will(
if row[6] != properties: if row[6] != properties:
raise ValueError("Invalid properties %s / %s" % (row[6], properties)) raise ValueError("Invalid properties %s / %s" % (row[6], properties))
except ValueError as err:
raise ValueError(str(err) + f" at index {idx}") from err
finally:
con.close()
return row[0] return row[0]