Files

214 lines
7.8 KiB
Python

from datetime import timedelta
from pathlib import Path
from types import SimpleNamespace
import shutil
import subprocess
import os
import pytest
from btcdata_archive.archive import run
from btcdata_archive import sqlite_archive
from btcdata_archive.schema import TABLES
from btcdata_archive.weeks import timestamp
from test_contract import START
class PG:
def __init__(self, config):
self.config = config
self.data = {START: [(1, '1.00000000', timestamp(START))], START + timedelta(days=14): [(3, '3.00000000', timestamp(START + timedelta(days=14)))]}
self.states = {}
self.events = []
self.locked = False
self.change = False
def lock(self):
self.locked = True
def alive(self):
if not self.locked:
raise RuntimeError('session lost')
def check_schema(self):
pass
def first_week(self, cutoff):
return min([s for s, data in self.data.items() if data and s < cutoff] + [s for s, state in self.states.items() if state != 'complete'], default=None)
def counts(self, start, end):
return {t.name: len(self.data.get(start, [])) if t == TABLES[0] else 0 for t in TABLES}
def archived_version(self, start):
return None
def estimated_bytes(self, start):
return 1024
def ledger(self, start, status, *args, **kwargs):
self.alive()
self.events.append((start, status))
self.states[start] = status
def export(self, path, start):
self.alive()
db = sqlite_archive.connect(path)
sqlite_archive.upsert(db, TABLES[0], self.data.get(start, []))
db.commit()
db.close()
def prune(self, path, start, cutoff, progress):
self.alive()
if self.change:
self.data[start][0] = (1, '9.00000000', timestamp(start))
self.change = False
db = sqlite_archive.connect(path, readonly=True)
archived = set(db.execute('SELECT * FROM binance_price'))
db.close()
old = self.data.get(start, [])
self.data[start] = [r for r in old if r not in archived]
progress({'binance_price': len(old) - len(self.data[start])})
if self.data[start]:
raise RuntimeError('reconciliation required')
return {'binance_price': len(old)}
class Store:
def __init__(self):
self.objects = {}
self.fail = None
self.pg = None
def require_versioning(self):
if self.fail == 'versioning':
raise RuntimeError('disabled')
def head(self, key):
if key not in self.objects:
return None
version, data, checksum = self.objects[key][-1]
return {'VersionId': version, 'ContentLength': len(data), 'Metadata': {'sha256': checksum}}
def upload(self, key, path, checksum):
if self.fail == 'upload':
raise OSError('sensitive password')
version = str(len(self.objects.get(key, [])) + 1)
self.objects.setdefault(key, []).append((version, path.read_bytes(), checksum))
if self.fail == 'lock_loss':
self.pg.locked = False
return version
def download(self, key, version, path):
if self.fail == 'download':
raise OSError('sensitive payload')
data = next(data for v, data, sha in self.objects[key] if v == version)
path.write_bytes(data if self.fail != 'checksum' else data + b'bad')
@pytest.fixture
def system(tmp_path):
config = SimpleNamespace(prefix='btcprod/weeks/', work_dir=tmp_path, pg={'dbname': 'btcprod'}, batch_rows=2)
pg, store = PG(config), Store()
store.pg = pg
return config, pg, store
def test_sequential_gap_and_fixed_cutoff(system):
config, pg, store = system
cutoff = START + timedelta(days=21)
run(pg, store, config, cutoff, 'archive-and-prune')
assert [s for s, state in pg.events if state == 'exporting'] == [START + timedelta(days=d) for d in (0,7,14)]
assert len(store.objects) == 3
assert all(not values for values in pg.data.values())
for i, (s, state) in enumerate(pg.events):
if state == 'exporting' and s != START:
assert pg.events[i-1][1] == 'complete'
assert not list(config.work_dir.iterdir())
def test_verify_only_bounded_resume_preserves_old_rows(system, tmp_path):
config, pg, store = system
cutoff = START + timedelta(days=21)
run(pg, store, config, cutoff, 'verify-only', 1)
assert pg.states[START] == 'verified'
assert pg.data[START]
run(pg, store, config, cutoff, 'archive-and-prune', 1)
assert not pg.data[START]
pg.data[START] = [(2, '2.00000000', timestamp(START))]
run(pg, store, config, cutoff, 'archive-and-prune', 1)
key = config.prefix + START.date().isoformat() + '.sqlite'
path = tmp_path / 'restore.sqlite'
store.download(key, store.head(key)['VersionId'], path)
db = sqlite_archive.connect(path, readonly=True)
assert db.execute('SELECT id FROM binance_price ORDER BY id').fetchall() == [(1,), (2,)]
db.close()
assert len(store.objects[key]) == 3
@pytest.mark.parametrize('failure', ['upload', 'download', 'checksum', 'versioning', 'lock_loss'])
def test_failure_never_prunes_unverified_or_advances(system, failure, capsys):
config, pg, store = system
store.fail = failure
with pytest.raises(Exception):
run(pg, store, config, START + timedelta(days=21), 'archive-and-prune')
assert pg.data[START]
assert not any(s > START for s, state in pg.events)
assert 'sensitive' not in capsys.readouterr().out
assert not list(config.work_dir.iterdir())
def test_changed_row_reconciliation(system):
config, pg, store = system
pg.change = True
with pytest.raises(RuntimeError):
run(pg, store, config, START + timedelta(days=21), 'archive-and-prune')
assert pg.states[START] == 'failed'
assert pg.data[START][0][1] == '9.00000000'
run(pg, store, config, START + timedelta(days=21), 'archive-and-prune', 1)
assert not pg.data[START]
def test_full_disk(system, monkeypatch):
config, pg, store = system
monkeypatch.setattr(shutil, 'disk_usage', lambda p: SimpleNamespace(free=0))
with pytest.raises(RuntimeError, match='storage'):
run(pg, store, config, START + timedelta(days=21), 'archive-and-prune')
assert pg.data[START]
assert not store.objects
def test_missing_previously_verified_object(system, monkeypatch):
config, pg, store = system
monkeypatch.setattr(pg, 'archived_version', lambda start: 'previous-version')
with pytest.raises(RuntimeError, match='operator recovery'):
run(pg, store, config, START + timedelta(days=21), 'archive-and-prune')
assert pg.data[START]
assert not store.objects
def test_empty(system):
config, pg, store = system
pg.data.clear()
run(pg, store, config, START + timedelta(days=21), 'archive-and-prune')
assert not store.objects
def test_wrapper_interpreter_override_and_arguments(tmp_path):
interpreter = tmp_path / 'python-test'
interpreter.write_text('#!/bin/sh\nprintf "%s\\n" "$@"\nexit 17\n')
interpreter.chmod(0o755)
script = Path('scripts/archive-old-weeks.sh').resolve()
result = subprocess.run([str(script), '--mode', 'verify-only', '--max-weeks', '1'],
cwd=tmp_path, env={**os.environ, 'PYTHON': str(interpreter)},
capture_output=True, text=True)
assert result.returncode == 17
assert result.stdout.splitlines() == ['-m', 'btcdata_archive', 'run', '--mode', 'verify-only', '--max-weeks', '1']
def test_wrapper_requires_mode():
import sys
env = {**os.environ, 'PATH': str(Path(sys.executable).parent) + ':' + os.environ['PATH']}
result = subprocess.run(['scripts/archive-old-weeks.sh'], env=env, capture_output=True, text=True)
assert result.returncode == 2
assert '--mode' in result.stderr