100 lines
4.2 KiB
Python
100 lines
4.2 KiB
Python
import io
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from botocore.response import StreamingBody
|
|
from botocore.stub import Stubber
|
|
|
|
from btcdata_archive.errors import SafetyError
|
|
from btcdata_archive.object_store import ObjectStore
|
|
from btcdata_archive.postgres import Postgres
|
|
from btcdata_archive.schema import TABLES
|
|
|
|
|
|
def store():
|
|
return ObjectStore(SimpleNamespace(endpoint='http://127.0.0.1:19000', access_key='test', secret_key='test', region='us-east-1', bucket='test'))
|
|
|
|
|
|
def test_versioning_enforced():
|
|
backend = store()
|
|
with Stubber(backend.client) as stub:
|
|
stub.add_response('get_bucket_versioning', {'Status':'Suspended'}, {'Bucket':'test'})
|
|
with pytest.raises(SafetyError, match='versioning'):
|
|
backend.require_versioning()
|
|
|
|
|
|
def test_head_missing_not_forbidden():
|
|
backend = store()
|
|
with Stubber(backend.client) as stub:
|
|
stub.add_client_error('head_object', 'NoSuchKey', http_status_code=404, expected_params={'Bucket':'test','Key':'week.sqlite'})
|
|
assert backend.head('week.sqlite') is None
|
|
stub.add_client_error('head_object', 'AccessDenied', http_status_code=403, expected_params={'Bucket':'test','Key':'week.sqlite'})
|
|
with pytest.raises(Exception):
|
|
backend.head('week.sqlite')
|
|
|
|
|
|
def test_download_exact_version(tmp_path):
|
|
backend = store()
|
|
body = StreamingBody(io.BytesIO(b'archive'), 7)
|
|
with Stubber(backend.client) as stub:
|
|
stub.add_response('get_object', {'VersionId':'v1','Body':body}, {'Bucket':'test','Key':'week.sqlite','VersionId':'v1'})
|
|
path = tmp_path / 'archive'
|
|
backend.download('week.sqlite','v1',path)
|
|
assert path.read_bytes() == b'archive'
|
|
stub.add_response('get_object', {'VersionId':'v2','Body':StreamingBody(io.BytesIO(b'bad'), 3)}, {'Bucket':'test','Key':'week.sqlite','VersionId':'v1'})
|
|
with pytest.raises(SafetyError, match='wrong version'):
|
|
backend.download('week.sqlite','v1',path)
|
|
|
|
|
|
def test_upload_captures_completion_version(tmp_path):
|
|
backend = store()
|
|
path = tmp_path / 'archive'
|
|
path.write_bytes(b'archive')
|
|
with Stubber(backend.client) as stub:
|
|
stub.add_response('get_bucket_versioning', {'Status':'Enabled'}, {'Bucket':'test'})
|
|
stub.add_response('create_multipart_upload', {'UploadId':'upload'}, {'Bucket':'test','Key':'week.sqlite','ContentType':'application/vnd.sqlite3','Metadata':{'sha256':'digest'}})
|
|
stub.add_response('upload_part', {'ETag':'not-a-checksum'}, {'Bucket':'test','Key':'week.sqlite','UploadId':'upload','PartNumber':1,'Body':b'archive'})
|
|
stub.add_response('complete_multipart_upload', {'VersionId':'verified-version'}, {'Bucket':'test','Key':'week.sqlite','UploadId':'upload','MultipartUpload':{'Parts':[{'PartNumber':1,'ETag':'not-a-checksum'}]}})
|
|
assert backend.upload('week.sqlite',path,'digest') == 'verified-version'
|
|
|
|
|
|
class Result:
|
|
def __init__(self, rows):
|
|
self.rows = rows
|
|
def fetchone(self):
|
|
return self.rows[0]
|
|
def fetchall(self):
|
|
return self.rows
|
|
|
|
|
|
class SchemaConnection:
|
|
def __init__(self, change=None):
|
|
self.change = change
|
|
def execute(self, query, params):
|
|
table = next(t for t in TABLES if params[0] == 'public.' + t.name)
|
|
if 'relkind' in query:
|
|
return Result([('r',False)])
|
|
if 'pg_trigger' in query or 'confdeltype' in query:
|
|
return Result([(0,)])
|
|
if 'format_type' in query:
|
|
columns = list(table.columns)
|
|
if self.change == 'added':
|
|
columns.append(('extra','text',True))
|
|
elif self.change == 'removed':
|
|
columns.pop()
|
|
elif self.change == 'type':
|
|
columns[0] = ('id','integer',False)
|
|
return Result(columns)
|
|
return Result([([table.pk if self.change != 'pk' else 'wrong'],)])
|
|
|
|
|
|
@pytest.mark.parametrize('change', [None,'added','removed','type','pk'])
|
|
def test_source_schema_fail_closed(change):
|
|
pg = Postgres.__new__(Postgres)
|
|
pg.conn = SchemaConnection(change)
|
|
if change is None:
|
|
pg.check_schema()
|
|
else:
|
|
with pytest.raises(SafetyError, match='migration required'):
|
|
pg.check_schema()
|