sashatrask
648 lines
25 KiB
Python
648 lines
25 KiB
Python
import multiprocessing
|
|
import os
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
import sys
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from unittest import mock
|
|
import uuid
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / 'app'
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
import console_runner
|
|
from process_identity import current_process_identity
|
|
from result_spool import ResultSpool, SpoolCapacityError
|
|
from scanner_db import ScannerDB
|
|
|
|
|
|
class SharedPublisherDB:
|
|
def __init__(self, lease, entered, release_gate, ingest_count, hold=False, owner_state=True):
|
|
self.conn = SimpleNamespace(is_postgres=True)
|
|
self.lease = lease
|
|
self.entered = entered
|
|
self.release_gate = release_gate
|
|
self.ingest_count = ingest_count
|
|
self.hold = hold
|
|
self.owner_state = owner_state
|
|
self.owns_lease = False
|
|
|
|
def require_runtime_safety_schema(self):
|
|
return True
|
|
|
|
def try_acquire_result_spool_publisher(self):
|
|
self.owns_lease = self.lease.acquire(False)
|
|
return self.owns_lease
|
|
|
|
def result_spool_publisher_state(self):
|
|
if not self.owner_state:
|
|
return {
|
|
'status': 'held', 'authenticated': False,
|
|
'application_name': 'unknown', 'holder_identity': 'unknown',
|
|
}
|
|
return {
|
|
'status': 'held', 'authenticated': True,
|
|
'application_name': 'truf-source:live_owner', 'pid': 123,
|
|
'state': 'active', 'holder_identity': '123:fixture-start',
|
|
}
|
|
|
|
def release_result_spool_publisher(self):
|
|
if not self.owns_lease:
|
|
return False
|
|
self.owns_lease = False
|
|
self.lease.release()
|
|
return True
|
|
|
|
def ingest_scan_event(self, envelope):
|
|
with self.ingest_count.get_lock():
|
|
self.ingest_count.value += 1
|
|
if self.hold:
|
|
self.entered.set()
|
|
if not self.release_gate.wait(15):
|
|
raise RuntimeError('fixture publisher release timed out')
|
|
return {
|
|
'ingested': True,
|
|
'scan_event_id': envelope['scan_event_id'],
|
|
'scan_event_hash': envelope['scan_event_hash'],
|
|
'duplicate': False,
|
|
}
|
|
|
|
|
|
def publisher_worker(
|
|
directory,
|
|
lease,
|
|
entered,
|
|
release_gate,
|
|
ingest_count,
|
|
outcomes,
|
|
hold=False,
|
|
stop_event=None,
|
|
):
|
|
try:
|
|
spool = ResultSpool(
|
|
directory,
|
|
max_event_bytes=4096,
|
|
max_events=20,
|
|
max_total_bytes=65536,
|
|
min_free_bytes=0,
|
|
lock_timeout_sec=2,
|
|
)
|
|
db = SharedPublisherDB(lease, entered, release_gate, ingest_count, hold=hold)
|
|
drained = console_runner.wait_for_result_spool_ready(
|
|
spool,
|
|
db,
|
|
stop_event=stop_event,
|
|
wait_seconds=0.05,
|
|
diagnostic_interval=0.2,
|
|
)
|
|
outcomes.put(('ok', len(drained)))
|
|
except KeyboardInterrupt:
|
|
outcomes.put(('stopped', 0))
|
|
except BaseException as exc:
|
|
outcomes.put((type(exc).__name__, str(exc)))
|
|
|
|
|
|
class CapacityProgressDB:
|
|
conn = SimpleNamespace(is_postgres=True)
|
|
|
|
def require_runtime_safety_schema(self):
|
|
return True
|
|
|
|
def result_spool_reservation_progress(self, reservations):
|
|
return {
|
|
'safe_progress': bool(reservations),
|
|
'counts': {
|
|
'exact_live': len(reservations),
|
|
'exact_expired': 0,
|
|
'stale_or_reassigned': 0,
|
|
},
|
|
'earliest_progress_in_sec': 60,
|
|
}
|
|
|
|
|
|
def capacity_worker(directory, index, start, stop_event, outcomes):
|
|
try:
|
|
spool = ResultSpool(
|
|
directory,
|
|
max_event_bytes=1024,
|
|
max_events=100,
|
|
max_total_bytes=2500,
|
|
min_free_bytes=0,
|
|
lock_timeout_sec=2,
|
|
)
|
|
start.wait(10)
|
|
reservation = console_runner.reserve_result_spool_claims(
|
|
spool,
|
|
CapacityProgressDB(),
|
|
f'source-{index}:{os.getpid()}:cycle',
|
|
1,
|
|
60,
|
|
stop_event=stop_event,
|
|
wait_seconds=0.05,
|
|
diagnostic_interval=1,
|
|
)
|
|
time.sleep(0.05)
|
|
if not spool.release_reservation(reservation):
|
|
raise RuntimeError('worker reservation release failed')
|
|
outcomes.put(('acquired', index))
|
|
except KeyboardInterrupt:
|
|
outcomes.put(('stopped', index))
|
|
except BaseException as exc:
|
|
outcomes.put((type(exc).__name__, index))
|
|
|
|
|
|
class ResultSpoolBackpressureTests(unittest.TestCase):
|
|
@staticmethod
|
|
def context():
|
|
return multiprocessing.get_context('spawn')
|
|
|
|
def test_completion_handoff_cannot_be_misread_as_reservation_renewal_loss(self):
|
|
class RacingSpool:
|
|
def __init__(self):
|
|
self.write_entered = threading.Event()
|
|
self.allow_write = threading.Event()
|
|
self.renew_calls = 0
|
|
|
|
def write_event(self, *args, **kwargs):
|
|
self.write_entered.set()
|
|
if not self.allow_write.wait(5):
|
|
raise RuntimeError('fixture write release timed out')
|
|
return SimpleNamespace(event_id='durable-event')
|
|
|
|
def renew_reservation(self, reservation_id, lease_seconds):
|
|
self.renew_calls += 1
|
|
return False
|
|
|
|
spool = RacingSpool()
|
|
lease_token = 'lease-token'
|
|
active = {lease_token}
|
|
reserved = {lease_token}
|
|
active_lock = threading.Lock()
|
|
outcomes = {}
|
|
renewal_attempted = threading.Event()
|
|
|
|
writer = threading.Thread(target=lambda: outcomes.__setitem__('write',
|
|
console_runner.write_reserved_result_spool_event(
|
|
spool, {}, 'reservation', 7, lease_token,
|
|
reserved, active, active_lock,
|
|
).event_id
|
|
))
|
|
|
|
def renew():
|
|
renewal_attempted.set()
|
|
outcomes['renew'] = console_runner.renew_active_result_spool_reservation(
|
|
spool, 'reservation', 1800, reserved, active_lock,
|
|
)
|
|
|
|
writer.start()
|
|
self.assertTrue(spool.write_entered.wait(5))
|
|
heartbeat = threading.Thread(target=renew)
|
|
heartbeat.start()
|
|
self.assertTrue(renewal_attempted.wait(5))
|
|
spool.allow_write.set()
|
|
writer.join(5)
|
|
heartbeat.join(5)
|
|
self.assertFalse(writer.is_alive())
|
|
self.assertFalse(heartbeat.is_alive())
|
|
self.assertEqual(outcomes, {'write': 'durable-event', 'renew': True})
|
|
self.assertEqual(spool.renew_calls, 0)
|
|
self.assertEqual(active, set())
|
|
self.assertEqual(reserved, set())
|
|
|
|
def test_absent_reservation_with_active_token_still_fails_closed(self):
|
|
spool = SimpleNamespace(renew_reservation=mock.Mock(return_value=False))
|
|
reserved = {'lease-token'}
|
|
self.assertFalse(console_runner.renew_active_result_spool_reservation(
|
|
spool, 'reservation', 1800, reserved, threading.Lock(),
|
|
))
|
|
spool.renew_reservation.assert_called_once_with('reservation', 1800)
|
|
|
|
def test_two_publishers_wait_handoff_and_ingest_once(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
directory = os.path.join(temp_dir, 'spool')
|
|
spool = ResultSpool(
|
|
directory, max_event_bytes=4096, max_events=20,
|
|
max_total_bytes=65536, min_free_bytes=0,
|
|
)
|
|
reservation = spool.reserve_claims('unrelated-live-claim', 1, 60, bytes_per_event=1024)
|
|
event_id = str(uuid.uuid4())
|
|
spool.write_event({'version': 1, 'scan_event_id': event_id, 'result': {}})
|
|
|
|
context = self.context()
|
|
lease = context.Lock()
|
|
entered = context.Event()
|
|
release_gate = context.Event()
|
|
ingest_count = context.Value('i', 0)
|
|
outcomes = context.Queue()
|
|
owner = context.Process(
|
|
target=publisher_worker,
|
|
args=(directory, lease, entered, release_gate, ingest_count, outcomes, True),
|
|
)
|
|
owner.start()
|
|
self.assertTrue(entered.wait(10))
|
|
waiter = context.Process(
|
|
target=publisher_worker,
|
|
args=(directory, lease, entered, release_gate, ingest_count, outcomes),
|
|
)
|
|
waiter.start()
|
|
time.sleep(0.3)
|
|
release_gate.set()
|
|
for process in (owner, waiter):
|
|
process.join(20)
|
|
self.assertFalse(process.is_alive())
|
|
self.assertEqual(process.exitcode, 0)
|
|
|
|
results = sorted(outcomes.get(timeout=5) for _ in range(2))
|
|
self.assertEqual(results, [('ok', 0), ('ok', 1)])
|
|
self.assertEqual(ingest_count.value, 1)
|
|
self.assertFalse(os.path.exists(spool._event_path(event_id)))
|
|
self.assertTrue(os.path.exists(spool._reservation_path(reservation)))
|
|
self.assertTrue(spool.release_reservation(reservation))
|
|
|
|
def test_waiter_honors_shutdown_without_touching_pending_event(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
directory = os.path.join(temp_dir, 'spool')
|
|
spool = ResultSpool(
|
|
directory, max_event_bytes=4096, max_events=20,
|
|
max_total_bytes=65536, min_free_bytes=0,
|
|
)
|
|
event_id = str(uuid.uuid4())
|
|
spool.write_event({'version': 1, 'scan_event_id': event_id, 'result': {}})
|
|
|
|
context = self.context()
|
|
lease = context.Lock()
|
|
entered = context.Event()
|
|
release_gate = context.Event()
|
|
stop_event = context.Event()
|
|
ingest_count = context.Value('i', 0)
|
|
outcomes = context.Queue()
|
|
owner = context.Process(
|
|
target=publisher_worker,
|
|
args=(directory, lease, entered, release_gate, ingest_count, outcomes, True),
|
|
)
|
|
owner.start()
|
|
self.assertTrue(entered.wait(10))
|
|
waiter = context.Process(
|
|
target=publisher_worker,
|
|
args=(directory, lease, entered, release_gate, ingest_count, outcomes, False, stop_event),
|
|
)
|
|
waiter.start()
|
|
stop_event.set()
|
|
waiter.join(15)
|
|
self.assertFalse(waiter.is_alive())
|
|
self.assertEqual(waiter.exitcode, 0)
|
|
waiter_result = outcomes.get(timeout=5)
|
|
self.assertEqual(waiter_result, ('stopped', 0))
|
|
self.assertTrue(os.path.exists(spool._event_path(event_id)))
|
|
|
|
release_gate.set()
|
|
owner.join(20)
|
|
self.assertFalse(owner.is_alive())
|
|
self.assertEqual(owner.exitcode, 0)
|
|
self.assertEqual(outcomes.get(timeout=5), ('ok', 1))
|
|
self.assertEqual(ingest_count.value, 1)
|
|
|
|
def test_unknown_competing_publisher_identity_is_fatal(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
directory = os.path.join(temp_dir, 'spool')
|
|
spool = ResultSpool(
|
|
directory, max_event_bytes=4096, max_events=20,
|
|
max_total_bytes=65536, min_free_bytes=0,
|
|
)
|
|
spool.write_event({'version': 1, 'scan_event_id': str(uuid.uuid4()), 'result': {}})
|
|
context = self.context()
|
|
lease = context.Lock()
|
|
self.assertTrue(lease.acquire(False))
|
|
try:
|
|
db = SharedPublisherDB(
|
|
lease, context.Event(), context.Event(), context.Value('i', 0),
|
|
owner_state=False,
|
|
)
|
|
with self.assertRaisesRegex(RuntimeError, 'unknown or unauthenticated'):
|
|
console_runner.wait_for_result_spool_ready(
|
|
spool, db, wait_seconds=0.01, diagnostic_interval=0.01,
|
|
)
|
|
finally:
|
|
lease.release()
|
|
|
|
@staticmethod
|
|
def empty_spool():
|
|
return SimpleNamespace(
|
|
next_pending_event=lambda: None,
|
|
assert_claims_allowed=lambda: True,
|
|
)
|
|
|
|
@staticmethod
|
|
def sequenced_db(acquisitions, states=None, release=True, state_error=None):
|
|
class DB:
|
|
conn = SimpleNamespace(is_postgres=True)
|
|
|
|
def __init__(self):
|
|
self.acquisitions = list(acquisitions)
|
|
self.states = list(states or [])
|
|
self.release_result = release
|
|
self.state_error = state_error
|
|
self.acquire_calls = 0
|
|
self.release_calls = 0
|
|
|
|
def require_runtime_safety_schema(self):
|
|
return True
|
|
|
|
def try_acquire_result_spool_publisher(self):
|
|
self.acquire_calls += 1
|
|
return self.acquisitions.pop(0)
|
|
|
|
def result_spool_publisher_state(self):
|
|
if self.state_error:
|
|
raise self.state_error
|
|
return self.states.pop(0)
|
|
|
|
def release_result_spool_publisher(self):
|
|
self.release_calls += 1
|
|
return self.release_result
|
|
|
|
return DB()
|
|
|
|
def test_release_between_try_and_inspect_retries_acquire(self):
|
|
db = self.sequenced_db(
|
|
[False, True],
|
|
states=[{'status': 'free', 'authenticated': True}],
|
|
)
|
|
with mock.patch.object(console_runner.time, 'sleep') as sleep:
|
|
self.assertEqual(
|
|
console_runner.wait_for_result_spool_ready(
|
|
self.empty_spool(), db, wait_seconds=0.1, diagnostic_interval=30,
|
|
),
|
|
{},
|
|
)
|
|
self.assertEqual(db.acquire_calls, 2)
|
|
self.assertEqual(db.release_calls, 1)
|
|
sleep.assert_called_once_with(0.1)
|
|
|
|
def test_holder_handoff_between_observations_retries_safely(self):
|
|
owner_a = {
|
|
'status': 'held', 'authenticated': True,
|
|
'application_name': 'truf-source:github', 'holder_identity': '10:start-a',
|
|
}
|
|
owner_b = {
|
|
'status': 'held', 'authenticated': True,
|
|
'application_name': 'truf-source:gitlab', 'holder_identity': '11:start-b',
|
|
}
|
|
db = self.sequenced_db([False, False, True], states=[owner_a, owner_b])
|
|
with mock.patch.object(console_runner.time, 'sleep'):
|
|
self.assertEqual(
|
|
console_runner.wait_for_result_spool_ready(
|
|
self.empty_spool(), db, wait_seconds=0.1, diagnostic_interval=30,
|
|
),
|
|
{},
|
|
)
|
|
self.assertEqual(db.acquire_calls, 3)
|
|
self.assertEqual(db.release_calls, 1)
|
|
|
|
def test_owner_lookup_error_remains_fatal(self):
|
|
db = self.sequenced_db([False], state_error=OSError('catalog lookup failed'))
|
|
with self.assertRaisesRegex(OSError, 'catalog lookup failed'):
|
|
console_runner.wait_for_result_spool_ready(
|
|
self.empty_spool(), db, wait_seconds=0.1, diagnostic_interval=30,
|
|
)
|
|
self.assertEqual(db.acquire_calls, 1)
|
|
self.assertEqual(db.release_calls, 0)
|
|
|
|
def test_shutdown_during_release_inspection_race_stops_wait(self):
|
|
class StopDuringWait:
|
|
def is_set(self):
|
|
return False
|
|
|
|
def wait(self, timeout):
|
|
return True
|
|
|
|
db = self.sequenced_db(
|
|
[False], states=[{'status': 'free', 'authenticated': True}],
|
|
)
|
|
with self.assertRaises(KeyboardInterrupt):
|
|
console_runner.wait_for_result_spool_ready(
|
|
self.empty_spool(), db, stop_event=StopDuringWait(),
|
|
wait_seconds=0.1, diagnostic_interval=30,
|
|
)
|
|
self.assertEqual(db.acquire_calls, 1)
|
|
self.assertEqual(db.release_calls, 0)
|
|
|
|
def test_exact_release_is_required_on_acquired_session(self):
|
|
successful = self.sequenced_db([True], release=True)
|
|
self.assertEqual(
|
|
console_runner.drain_result_spool(self.empty_spool(), successful), {},
|
|
)
|
|
self.assertEqual((successful.acquire_calls, successful.release_calls), (1, 1))
|
|
|
|
uncertain = self.sequenced_db([True], release=False)
|
|
with self.assertRaisesRegex(RuntimeError, 'lease release was not confirmed'):
|
|
console_runner.drain_result_spool(self.empty_spool(), uncertain)
|
|
self.assertEqual((uncertain.acquire_calls, uncertain.release_calls), (1, 1))
|
|
|
|
def test_scanner_db_tracks_exact_same_session_acquire_release(self):
|
|
class Cursor:
|
|
def __init__(self, row):
|
|
self.row = row
|
|
|
|
def fetchone(self):
|
|
return self.row
|
|
|
|
class Connection:
|
|
is_postgres = True
|
|
|
|
def __init__(self):
|
|
self.closed = False
|
|
self.calls = []
|
|
|
|
def execute(self, sql, params=None):
|
|
self.calls.append(sql)
|
|
if 'pg_try_advisory_lock' in sql:
|
|
return Cursor({'acquired': True})
|
|
if 'pg_advisory_unlock' in sql:
|
|
return Cursor({'released': True})
|
|
raise AssertionError(sql)
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
connection = Connection()
|
|
db = ScannerDB.__new__(ScannerDB)
|
|
db.conn = connection
|
|
db._result_spool_publisher_held = False
|
|
db.last_error = ''
|
|
self.assertTrue(db.try_acquire_result_spool_publisher())
|
|
with self.assertRaisesRegex(RuntimeError, 'already held'):
|
|
db.try_acquire_result_spool_publisher()
|
|
self.assertTrue(db.release_result_spool_publisher())
|
|
self.assertFalse(db.release_result_spool_publisher())
|
|
self.assertEqual(sum('pg_try_advisory_lock' in sql for sql in connection.calls), 1)
|
|
self.assertEqual(sum('pg_advisory_unlock' in sql for sql in connection.calls), 1)
|
|
|
|
def test_uncertain_release_closes_session_to_prevent_lock_leak(self):
|
|
class Cursor:
|
|
def fetchone(self):
|
|
return {'released': False}
|
|
|
|
class Connection:
|
|
is_postgres = True
|
|
|
|
def __init__(self):
|
|
self.closed = False
|
|
|
|
def execute(self, sql, params=None):
|
|
return Cursor()
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def close(self):
|
|
self.closed = True
|
|
|
|
connection = Connection()
|
|
db = ScannerDB.__new__(ScannerDB)
|
|
db.conn = connection
|
|
db._result_spool_publisher_held = True
|
|
db.last_error = ''
|
|
self.assertFalse(db.release_result_spool_publisher())
|
|
self.assertTrue(connection.closed)
|
|
self.assertIsNone(db.conn)
|
|
self.assertFalse(db._result_spool_publisher_held)
|
|
|
|
def test_twelve_sources_wait_for_capacity_then_handoff_without_loss(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
directory = os.path.join(temp_dir, 'spool')
|
|
spool = ResultSpool(
|
|
directory, max_event_bytes=1024, max_events=100,
|
|
max_total_bytes=2500, min_free_bytes=0,
|
|
)
|
|
blocker = spool.reserve_claims('unrelated-owner', 1, 60, bytes_per_event=1024)
|
|
context = self.context()
|
|
start = context.Event()
|
|
never_stop = context.Event()
|
|
outcomes = context.Queue()
|
|
processes = [
|
|
context.Process(
|
|
target=capacity_worker,
|
|
args=(directory, index, start, never_stop, outcomes),
|
|
)
|
|
for index in range(12)
|
|
]
|
|
for process in processes:
|
|
process.start()
|
|
start.set()
|
|
time.sleep(0.5)
|
|
self.assertTrue(os.path.exists(spool._reservation_path(blocker)))
|
|
self.assertTrue(spool.release_reservation(blocker))
|
|
for process in processes:
|
|
process.join(30)
|
|
self.assertFalse(process.is_alive())
|
|
self.assertEqual(process.exitcode, 0)
|
|
results = sorted(outcomes.get(timeout=5) for _ in processes)
|
|
self.assertEqual(results, [('acquired', index) for index in range(12)])
|
|
self.assertEqual(list(os.scandir(spool.reservation_directory)), [])
|
|
self.assertEqual(spool._pending_paths()[0], [])
|
|
|
|
def test_capacity_wait_honors_shutdown_and_retains_owner_reservation(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
directory = os.path.join(temp_dir, 'spool')
|
|
spool = ResultSpool(
|
|
directory, max_event_bytes=1024, max_events=100,
|
|
max_total_bytes=2500, min_free_bytes=0,
|
|
)
|
|
blocker = spool.reserve_claims('unrelated-owner', 1, 60, bytes_per_event=1024)
|
|
context = self.context()
|
|
start = context.Event()
|
|
stop_event = context.Event()
|
|
outcomes = context.Queue()
|
|
process = context.Process(
|
|
target=capacity_worker,
|
|
args=(directory, 1, start, stop_event, outcomes),
|
|
)
|
|
process.start()
|
|
start.set()
|
|
time.sleep(0.3)
|
|
stop_event.set()
|
|
process.join(15)
|
|
self.assertFalse(process.is_alive())
|
|
self.assertEqual(process.exitcode, 0)
|
|
self.assertEqual(outcomes.get(timeout=5), ('stopped', 1))
|
|
self.assertTrue(os.path.exists(spool._reservation_path(blocker)))
|
|
self.assertTrue(spool.release_reservation(blocker))
|
|
|
|
def test_intrinsically_impossible_capacity_request_remains_fatal(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
spool = ResultSpool(
|
|
os.path.join(temp_dir, 'spool'),
|
|
max_event_bytes=1024, max_events=100,
|
|
max_total_bytes=2500, min_free_bytes=0,
|
|
)
|
|
with self.assertRaises(SpoolCapacityError):
|
|
console_runner.reserve_result_spool_claims(
|
|
spool, CapacityProgressDB(), 'owner', 3, 60,
|
|
wait_seconds=0.01, diagnostic_interval=0.01,
|
|
)
|
|
|
|
def test_capacity_without_safe_progress_evidence_remains_fatal(self):
|
|
class UnsafeProgressDB(CapacityProgressDB):
|
|
def result_spool_reservation_progress(self, reservations):
|
|
return {'safe_progress': False}
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
spool = ResultSpool(
|
|
os.path.join(temp_dir, 'spool'),
|
|
max_event_bytes=1024, max_events=100,
|
|
max_total_bytes=2500, min_free_bytes=0,
|
|
)
|
|
blocker = spool.reserve_claims('unrelated-owner', 1, 60, bytes_per_event=1024)
|
|
with self.assertRaisesRegex(RuntimeError, 'unsafe or indeterminate'):
|
|
console_runner.reserve_result_spool_claims(
|
|
spool, UnsafeProgressDB(), 'owner', 1, 60,
|
|
wait_seconds=0.01, diagnostic_interval=0.01,
|
|
)
|
|
self.assertTrue(os.path.exists(spool._reservation_path(blocker)))
|
|
self.assertTrue(spool.release_reservation(blocker))
|
|
|
|
def test_exact_stopped_owner_refund_releases_only_its_unused_reservation(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict(os.environ, {
|
|
'SCANNER_DB_URL': '', 'DATABASE_URL': '',
|
|
}):
|
|
spool = ResultSpool(
|
|
os.path.join(temp_dir, 'spool'),
|
|
max_event_bytes=1024, max_events=20,
|
|
max_total_bytes=65536, min_free_bytes=0,
|
|
)
|
|
db = ScannerDB(db_path=os.path.join(temp_dir, 'scanner.db'))
|
|
identity = current_process_identity()
|
|
owner = f'source:{identity.pid}:cycle'
|
|
owned = spool.reserve_claims(owner, 1, 60, bytes_per_event=1024)
|
|
unrelated = spool.reserve_claims('other-owner', 1, 60, bytes_per_event=1024)
|
|
target = 'https://github.com/example/stopped-owner'
|
|
db.enqueue_targets('source', 'github', 'q', [target])
|
|
claims = db.claim_targets(
|
|
'source', 'github', 1, owner, 60,
|
|
return_rows=True, claim_batch=owned,
|
|
)
|
|
spool.bind_claims(owned, claims)
|
|
candidates = spool.stopped_owner_unused_reservations('source', identity)
|
|
self.assertEqual([row['reservation_id'] for row in candidates], [owned])
|
|
self.assertTrue(db.refund_stopped_result_spool_claims(candidates[0]))
|
|
self.assertTrue(spool.release_stopped_owner_reservation(owned, 'source', identity))
|
|
row = db.target_queue_item('source', 'github', target)
|
|
self.assertEqual((row['status'], row['attempts']), ('pending', 0))
|
|
self.assertFalse(os.path.exists(spool._reservation_path(owned)))
|
|
self.assertTrue(os.path.exists(spool._reservation_path(unrelated)))
|
|
self.assertTrue(spool.release_reservation(unrelated))
|
|
db.close()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|