77 lines
4.0 KiB
Python
77 lines
4.0 KiB
Python
"""Opt-in integration infrastructure. Always creates NEW disposable resources."""
|
|
import os
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
import uuid
|
|
|
|
import psycopg
|
|
from psycopg import sql
|
|
import pytest
|
|
|
|
from btcdata_archive.postgres import Postgres
|
|
from btcdata_archive.object_store import ObjectStore
|
|
from btcdata_archive.schema import TABLES
|
|
|
|
|
|
@pytest.fixture
|
|
def postgres_infra(tmp_path):
|
|
if not os.environ.get('TEST_PGHOST'):
|
|
pytest.skip('set TEST_PGHOST for disposable PostgreSQL infrastructure')
|
|
suffix = uuid.uuid4().hex
|
|
database, role, password = 'archive_test_' + suffix, 'archive_test_' + suffix, uuid.uuid4().hex
|
|
admin_args = {'host': os.environ['TEST_PGHOST'], 'port': int(os.environ.get('TEST_PGPORT', '5432')),
|
|
'user': os.environ.get('TEST_PGUSER', 'postgres'), 'password': os.environ.get('TEST_PGPASSWORD', 'disposable-test'),
|
|
'dbname': 'postgres'}
|
|
config = SimpleNamespace(pg={**admin_args, 'dbname': database, 'user': role, 'password': password},
|
|
endpoint=os.environ.get('TEST_S3_ENDPOINT_URL'), access_key=os.environ.get('TEST_S3_ACCESS_KEY', 'disposable-test'),
|
|
secret_key=os.environ.get('TEST_S3_SECRET_KEY', 'disposable-test-secret'),
|
|
region='us-east-1', bucket='archive-test-' + suffix, prefix='btcprod/weeks/', work_dir=tmp_path, batch_rows=2)
|
|
root = psycopg.connect(**admin_args, autocommit=True)
|
|
admin = pg = None
|
|
try:
|
|
root.execute(sql.SQL('CREATE DATABASE {}').format(sql.Identifier(database)))
|
|
root.execute(sql.SQL('CREATE ROLE {} LOGIN PASSWORD {}').format(sql.Identifier(role), sql.Literal(password)))
|
|
admin = psycopg.connect(**{**admin_args, 'dbname': database}, autocommit=True)
|
|
for table in TABLES:
|
|
cols = sql.SQL(',').join(sql.SQL('{} {} {}').format(sql.Identifier(name), sql.SQL(typ),
|
|
sql.SQL('PRIMARY KEY' if name == table.pk else '' if nullable else 'NOT NULL'))
|
|
for name, typ, nullable in table.columns)
|
|
admin.execute(sql.SQL('CREATE TABLE public.{} ({})').format(sql.Identifier(table.name), cols))
|
|
setup = (Path(__file__).resolve().parents[2] / 'sql/setup.sql').read_text()
|
|
setup = setup.replace(':"DBNAME"', sql.Identifier(database).as_string()).replace(':"runtime_role"', sql.Identifier(role).as_string())
|
|
admin.execute(setup)
|
|
pg = Postgres(config)
|
|
yield SimpleNamespace(config=config, pg=pg, admin=admin)
|
|
finally:
|
|
if pg:
|
|
pg.close()
|
|
if admin:
|
|
admin.close()
|
|
root.execute(sql.SQL('DROP DATABASE IF EXISTS {} WITH (FORCE)').format(sql.Identifier(database)))
|
|
root.execute(sql.SQL('DROP ROLE IF EXISTS {}').format(sql.Identifier(role)))
|
|
root.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def infra(postgres_infra):
|
|
if not os.environ.get('TEST_S3_ENDPOINT_URL'):
|
|
pytest.skip('set TEST_S3_ENDPOINT_URL for disposable RustFS infrastructure')
|
|
config = postgres_infra.config
|
|
store = ObjectStore(config)
|
|
bucket_created = False
|
|
instance = SimpleNamespace(config=config, pg=postgres_infra.pg, store=store, admin=postgres_infra.admin)
|
|
try:
|
|
store.client.create_bucket(Bucket=config.bucket)
|
|
bucket_created = True
|
|
store.client.put_bucket_versioning(Bucket=config.bucket, VersioningConfiguration={'Status': 'Enabled'})
|
|
yield instance
|
|
finally:
|
|
instance.pg.close()
|
|
if bucket_created:
|
|
# Test-only bucket with random unique name. Never use runtime cleanup in production.
|
|
for page in store.client.get_paginator('list_object_versions').paginate(Bucket=config.bucket):
|
|
items = [{'Key': v['Key'], 'VersionId': v['VersionId']} for v in page.get('Versions', []) + page.get('DeleteMarkers', [])]
|
|
if items:
|
|
store.client.delete_objects(Bucket=config.bucket, Delete={'Objects': items})
|
|
store.client.delete_bucket(Bucket=config.bucket)
|