mirror of
https://github.com/gadgethd/ukmesh.git
synced 2026-09-02 21:24:08 +00:00
374 lines
13 KiB
Python
374 lines
13 KiB
Python
import os
|
|
import pathlib
|
|
import sys
|
|
import time
|
|
import unittest
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
|
|
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
|
|
import worker # noqa: E402
|
|
|
|
|
|
class CursorTests(unittest.TestCase):
|
|
def test_more_than_one_batch_at_same_timestamp_is_total_and_stable(self):
|
|
observed_at = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
|
rows = [
|
|
{
|
|
'observed_at': observed_at,
|
|
'packet_hash': f'packet-{index:05d}',
|
|
'network': 'ukmesh',
|
|
'rx_node_id': 'receiver',
|
|
'topic': 'mesh/packet',
|
|
'raw_hex': f'{index:08x}',
|
|
}
|
|
for index in range(worker.GOLD_BATCH + 137)
|
|
]
|
|
cursors = sorted(worker.cursor_from_row(row) for row in rows)
|
|
|
|
checkpoint = worker.ExtractionCursor(
|
|
datetime(1970, 1, 1, tzinfo=timezone.utc),
|
|
'',
|
|
'',
|
|
'',
|
|
'',
|
|
'',
|
|
)
|
|
processed = []
|
|
while True:
|
|
batch = [cursor for cursor in cursors if cursor > checkpoint][
|
|
:worker.GOLD_BATCH
|
|
]
|
|
if not batch:
|
|
break
|
|
processed.extend(batch)
|
|
checkpoint = batch[-1]
|
|
|
|
self.assertEqual(len(rows), len(processed))
|
|
self.assertEqual(len(processed), len(set(processed)))
|
|
self.assertEqual(cursors, processed)
|
|
|
|
|
|
@unittest.skipUnless(
|
|
os.environ.get('TEST_DATABASE_URL'),
|
|
'TEST_DATABASE_URL is required for PostgreSQL integration tests',
|
|
)
|
|
class PostgreSQLLearnerTests(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.database_url = os.environ['TEST_DATABASE_URL']
|
|
|
|
def setUp(self):
|
|
self.schema = f'ml_test_{uuid.uuid4().hex}'
|
|
self.db = self._connect()
|
|
with self.db.cursor() as cur:
|
|
cur.execute(f'CREATE SCHEMA "{self.schema}"')
|
|
cur.execute(f'SET search_path TO "{self.schema}"')
|
|
cur.execute(
|
|
"""
|
|
CREATE TABLE ml_learner_state (
|
|
singleton BOOLEAN PRIMARY KEY DEFAULT TRUE CHECK (singleton),
|
|
cursor_observed_at TIMESTAMPTZ NOT NULL
|
|
DEFAULT '1970-01-01 00:00:00+00',
|
|
cursor_packet_hash TEXT NOT NULL DEFAULT '',
|
|
cursor_network TEXT NOT NULL DEFAULT '',
|
|
cursor_rx_node_id TEXT NOT NULL DEFAULT '',
|
|
cursor_topic TEXT NOT NULL DEFAULT '',
|
|
cursor_raw_hex TEXT NOT NULL DEFAULT '',
|
|
leader_token TEXT,
|
|
lease_expires_at TIMESTAMPTZ,
|
|
heartbeat_at TIMESTAMPTZ,
|
|
run_started_at TIMESTAMPTZ,
|
|
run_deadline_at TIMESTAMPTZ,
|
|
next_run_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
last_trained_at TIMESTAMPTZ,
|
|
next_training_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
model_version TEXT NOT NULL DEFAULT 'lightgbm-path-v1',
|
|
data_version TEXT NOT NULL DEFAULT 'gold-multibyte-v2',
|
|
last_terminal_reason TEXT,
|
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
|
);
|
|
INSERT INTO ml_learner_state (singleton) VALUES (TRUE);
|
|
|
|
CREATE TABLE nodes (
|
|
node_id TEXT NOT NULL,
|
|
network TEXT NOT NULL,
|
|
lat DOUBLE PRECISION,
|
|
lon DOUBLE PRECISION,
|
|
elevation_m DOUBLE PRECISION,
|
|
last_seen TIMESTAMPTZ,
|
|
iata TEXT
|
|
);
|
|
|
|
CREATE TABLE packets (
|
|
time TIMESTAMPTZ NOT NULL,
|
|
packet_hash TEXT NOT NULL,
|
|
network TEXT NOT NULL,
|
|
rx_node_id TEXT,
|
|
topic TEXT NOT NULL,
|
|
raw_hex TEXT,
|
|
path_hashes TEXT[],
|
|
path_hash_size_bytes INTEGER
|
|
);
|
|
|
|
CREATE TABLE ml_gold_paths (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
packet_hash TEXT NOT NULL,
|
|
network TEXT NOT NULL,
|
|
observed_at TIMESTAMPTZ NOT NULL,
|
|
hop_position INTEGER NOT NULL,
|
|
true_node_id TEXT NOT NULL,
|
|
hash_2char TEXT NOT NULL,
|
|
hash_4char TEXT,
|
|
hash_6char TEXT,
|
|
path_hash_size_bytes INTEGER,
|
|
observer_ids TEXT[],
|
|
rx_region TEXT,
|
|
UNIQUE (packet_hash, hop_position, true_node_id)
|
|
);
|
|
"""
|
|
)
|
|
self.db.commit()
|
|
|
|
def tearDown(self):
|
|
try:
|
|
self.db.rollback()
|
|
except Exception:
|
|
pass
|
|
self.db.close()
|
|
admin = self._connect(search_path=False)
|
|
with admin.cursor() as cur:
|
|
cur.execute(f'DROP SCHEMA IF EXISTS "{self.schema}" CASCADE')
|
|
admin.commit()
|
|
admin.close()
|
|
|
|
def _connect(self, search_path=True):
|
|
connection = psycopg2.connect(
|
|
self.database_url,
|
|
cursor_factory=psycopg2.extras.RealDictCursor,
|
|
)
|
|
connection.autocommit = False
|
|
if search_path and hasattr(self, 'schema'):
|
|
with connection.cursor() as cur:
|
|
cur.execute(f'SET search_path TO "{self.schema}"')
|
|
connection.commit()
|
|
return connection
|
|
|
|
def _claim(self, db):
|
|
original_url = worker.DATABASE_URL
|
|
original_lease = worker.ML_LEASE_SECONDS
|
|
worker.DATABASE_URL = self.database_url
|
|
worker.ML_LEASE_SECONDS = 600
|
|
try:
|
|
return worker.claim_leadership(db)
|
|
finally:
|
|
worker.DATABASE_URL = original_url
|
|
worker.ML_LEASE_SECONDS = original_lease
|
|
|
|
def _expire_lease(self):
|
|
with self.db.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
UPDATE ml_learner_state
|
|
SET lease_expires_at = NOW() - INTERVAL '1 second',
|
|
next_run_at = NOW() - INTERVAL '1 second'
|
|
"""
|
|
)
|
|
self.db.commit()
|
|
|
|
def test_single_leader_and_expired_leader_reclamation(self):
|
|
first = self._claim(self.db)
|
|
self.assertIsNotNone(first)
|
|
|
|
contender_db = self._connect()
|
|
try:
|
|
self.assertIsNone(self._claim(contender_db))
|
|
self._expire_lease()
|
|
reclaimed = self._claim(contender_db)
|
|
self.assertIsNotNone(reclaimed)
|
|
self.assertNotEqual(first.token, reclaimed.token)
|
|
finally:
|
|
contender_db.close()
|
|
|
|
def test_failed_batch_does_not_advance_checkpoint_or_skip_later_rows(self):
|
|
observed_at = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
|
with self.db.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO nodes
|
|
(node_id, network, lat, lon, elevation_m, last_seen, iata)
|
|
VALUES
|
|
('AAAA0001', 'ukmesh', 51.0, -1.0, 100, NOW(), 'GB-S'),
|
|
('BBBB0001', 'ukmesh', 51.1, -1.1, 120, NOW(), 'GB-S')
|
|
"""
|
|
)
|
|
packet_rows = [
|
|
(
|
|
observed_at,
|
|
f'p{index:05d}',
|
|
'ukmesh',
|
|
'RX000001',
|
|
'mesh/packet',
|
|
f'{index:08x}',
|
|
['AAAA', 'BBBB'],
|
|
2,
|
|
)
|
|
for index in range(worker.GOLD_BATCH + 1)
|
|
]
|
|
psycopg2.extras.execute_values(
|
|
cur,
|
|
"""
|
|
INSERT INTO packets
|
|
(time, packet_hash, network, rx_node_id, topic, raw_hex,
|
|
path_hashes, path_hash_size_bytes)
|
|
VALUES %s
|
|
""",
|
|
packet_rows,
|
|
page_size=1000,
|
|
)
|
|
cur.execute(
|
|
"""
|
|
CREATE INDEX packets_cursor_test_idx
|
|
ON packets (
|
|
time, packet_hash, network, COALESCE(rx_node_id, ''),
|
|
topic, COALESCE(raw_hex, '')
|
|
)
|
|
"""
|
|
)
|
|
self.db.commit()
|
|
|
|
guard = self._claim(self.db)
|
|
self.assertIsNotNone(guard)
|
|
self.assertEqual(worker.GOLD_BATCH, worker.extract_gold_paths(self.db, guard))
|
|
|
|
checkpoint_after_first = worker.get_checkpoint(self.db)
|
|
self.db.rollback()
|
|
self.assertEqual('p04999', checkpoint_after_first.packet_hash)
|
|
|
|
with self.db.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
CREATE FUNCTION reject_test_packet() RETURNS trigger
|
|
LANGUAGE plpgsql AS $$
|
|
BEGIN
|
|
IF NEW.packet_hash = 'p05000' THEN
|
|
RAISE EXCEPTION 'injected row failure';
|
|
END IF;
|
|
RETURN NEW;
|
|
END
|
|
$$;
|
|
CREATE TRIGGER reject_test_packet_trigger
|
|
BEFORE INSERT ON ml_gold_paths
|
|
FOR EACH ROW EXECUTE FUNCTION reject_test_packet();
|
|
"""
|
|
)
|
|
self.db.commit()
|
|
|
|
with self.assertRaisesRegex(Exception, 'injected row failure'):
|
|
worker.extract_gold_paths(self.db, guard)
|
|
self.db.rollback()
|
|
checkpoint_after_failure = worker.get_checkpoint(self.db)
|
|
self.db.rollback()
|
|
self.assertEqual(checkpoint_after_first, checkpoint_after_failure)
|
|
|
|
self._expire_lease()
|
|
restarted_db = self._connect()
|
|
try:
|
|
restarted_guard = self._claim(restarted_db)
|
|
self.assertIsNotNone(restarted_guard)
|
|
with restarted_db.cursor() as cur:
|
|
cur.execute('DROP TRIGGER reject_test_packet_trigger ON ml_gold_paths')
|
|
cur.execute('DROP FUNCTION reject_test_packet()')
|
|
restarted_db.commit()
|
|
|
|
self.assertEqual(
|
|
1,
|
|
worker.extract_gold_paths(restarted_db, restarted_guard),
|
|
)
|
|
self.assertEqual(
|
|
0,
|
|
worker.extract_gold_paths(restarted_db, restarted_guard),
|
|
)
|
|
with restarted_db.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
SELECT COUNT(*) AS total_rows,
|
|
COUNT(DISTINCT (packet_hash, hop_position, true_node_id))
|
|
AS distinct_rows
|
|
FROM ml_gold_paths
|
|
"""
|
|
)
|
|
counts = cur.fetchone()
|
|
restarted_db.rollback()
|
|
self.assertEqual(2 * (worker.GOLD_BATCH + 1), counts['total_rows'])
|
|
self.assertEqual(counts['total_rows'], counts['distinct_rows'])
|
|
finally:
|
|
restarted_db.close()
|
|
|
|
def test_training_rows_are_bounded_in_sql_by_whole_packets(self):
|
|
with self.db.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO nodes
|
|
(node_id, network, lat, lon, elevation_m, last_seen, iata)
|
|
VALUES
|
|
('AA000001', 'ukmesh', 51.0, -1.0, 100, NOW(), 'GB-S'),
|
|
('AA000002', 'ukmesh', 51.1, -1.1, 120, NOW(), 'GB-S')
|
|
"""
|
|
)
|
|
gold_rows = []
|
|
for packet_index in range(100):
|
|
for hop_position, node_id in enumerate(('AA000001', 'AA000002')):
|
|
gold_rows.append(
|
|
(
|
|
f'train-{packet_index:04d}',
|
|
'ukmesh',
|
|
datetime(2026, 1, 1, tzinfo=timezone.utc),
|
|
hop_position,
|
|
node_id,
|
|
'AA',
|
|
'AA00',
|
|
'AA0000',
|
|
2,
|
|
['RX000001'],
|
|
'GB-S',
|
|
)
|
|
)
|
|
psycopg2.extras.execute_values(
|
|
cur,
|
|
"""
|
|
INSERT INTO ml_gold_paths
|
|
(packet_hash, network, observed_at, hop_position, true_node_id,
|
|
hash_2char, hash_4char, hash_6char, path_hash_size_bytes,
|
|
observer_ids, rx_region)
|
|
VALUES %s
|
|
""",
|
|
gold_rows,
|
|
page_size=500,
|
|
)
|
|
self.db.commit()
|
|
|
|
original_min = worker.MIN_GOLD_ROWS
|
|
original_max = worker.MAX_TRAINING_GOLD_ROWS
|
|
worker.MIN_GOLD_ROWS = 1
|
|
worker.MAX_TRAINING_GOLD_ROWS = 20
|
|
try:
|
|
guard = worker.RunGuard('unit-test', time.monotonic() + 60)
|
|
X, y, meta, gold_ids = worker.build_training_data(self.db, guard)
|
|
finally:
|
|
worker.MIN_GOLD_ROWS = original_min
|
|
worker.MAX_TRAINING_GOLD_ROWS = original_max
|
|
self.db.rollback()
|
|
|
|
self.assertIsNotNone(X)
|
|
self.assertEqual(20, len(set(gold_ids.tolist())))
|
|
self.assertEqual(len(X), len(y))
|
|
self.assertEqual(len(X), len(meta))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|