begriffs open source - pg_scribe/blob - tests/test_restore.sh

   1 #!/usr/bin/env bash
   2 #
   3 # Test suite for pg_scribe --restore command
   4 #
   5 # This test suite:
   6 # - Creates temporary test databases
   7 # - Tests various --restore scenarios
   8 # - Verifies expected outcomes
   9 # - Cleans up all resources
  10 #
  11
  12 set -euo pipefail
  13
  14 # Colors for test output
  15 RED='\033[0;31m'
  16 GREEN='\033[0;32m'
  17 YELLOW='\033[0;33m'
  18 BLUE='\033[0;34m'
  19 NC='\033[0m' # No Color
  20
  21 # Test configuration
  22 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
  23 PG_SCRIBE="$SCRIPT_DIR/scripts/pg_scribe"
  24 TEST_DIR="/tmp/pg_scribe_test_restore_$$"
  25 TEST_DB_PREFIX="pg_scribe_restore_test_$$"
  26 PGUSER="${PGUSER:-postgres}"
  27
  28 # Test counters
  29 TESTS_RUN=0
  30 TESTS_PASSED=0
  31 TESTS_FAILED=0
  32
  33 # Cleanup tracking
  34 DATABASES_TO_CLEANUP=()
  35 SLOTS_TO_CLEANUP=()
  36
  37 #
  38 # Logging functions
  39 #
  40
  41 log_test() {
  42     echo -e "${BLUE}TEST:${NC} $*"
  43 }
  44
  45 log_pass() {
  46     echo -e "${GREEN}PASS:${NC} $*"
  47     ((TESTS_PASSED++))
  48 }
  49
  50 log_fail() {
  51     echo -e "${RED}FAIL:${NC} $*"
  52     ((TESTS_FAILED++))
  53 }
  54
  55 log_info() {
  56     echo -e "${YELLOW}INFO:${NC} $*"
  57 }
  58
  59 #
  60 # Helper functions
  61 #
  62
  63 run_psql() {
  64     local dbname="$1"
  65     shift
  66     psql -U "$PGUSER" -d "$dbname" -tAq "$@"
  67 }
  68
  69 query_db() {
  70     local dbname="$1"
  71     local query="$2"
  72     run_psql "$dbname" -c "$query" 2>/dev/null || true
  73 }
  74
  75 create_test_db() {
  76     local dbname="$1"
  77     log_info "Creating test database: $dbname"
  78
  79     # Drop if exists
  80     psql -U "$PGUSER" -d postgres -c "DROP DATABASE IF EXISTS $dbname;" &>/dev/null || true
  81
  82     # Create database
  83     psql -U "$PGUSER" -d postgres -c "CREATE DATABASE $dbname;" &>/dev/null
  84
  85     DATABASES_TO_CLEANUP+=("$dbname")
  86 }
  87
  88 # shellcheck disable=SC2317  # Function called from cleanup trap handler
  89 drop_test_db() {
  90     local dbname="$1"
  91     log_info "Dropping test database: $dbname"
  92
  93     # Terminate connections
  94     psql -U "$PGUSER" -d postgres -c "
  95         SELECT pg_terminate_backend(pid)
  96         FROM pg_stat_activity
  97         WHERE datname = '$dbname' AND pid <> pg_backend_pid();
  98     " &>/dev/null || true
  99
 100     # Drop database
 101     psql -U "$PGUSER" -d postgres -c "DROP DATABASE IF EXISTS $dbname;" &>/dev/null || true
 102 }
 103
 104 # shellcheck disable=SC2317  # Function called from cleanup trap handler
 105 drop_replication_slot() {
 106     local dbname="$1"
 107     local slot="$2"
 108     log_info "Dropping replication slot: $slot"
 109
 110     # Check if slot exists
 111     local exists
 112     exists=$(query_db "$dbname" "
 113         SELECT 1 FROM pg_replication_slots WHERE slot_name = '$slot';
 114     ")
 115
 116     if [[ -n "$exists" ]]; then
 117         # Drop slot
 118         query_db "$dbname" "SELECT pg_drop_replication_slot('$slot');" || true
 119     fi
 120 }
 121
 122 create_table_with_pk() {
 123     local dbname="$1"
 124     local table="$2"
 125     query_db "$dbname" "
 126         CREATE TABLE $table (
 127             id SERIAL PRIMARY KEY,
 128             name TEXT,
 129             created_at TIMESTAMP DEFAULT now()
 130         );
 131     "
 132 }
 133
 134 initialize_backup() {
 135     local dbname="$1"
 136     local backup_dir="$2"
 137     local slot="$3"
 138
 139     # Init backup system
 140     "$PG_SCRIBE" --init -d "$dbname" -f "$backup_dir" -S "$slot" -U "$PGUSER" &>/dev/null
 141     SLOTS_TO_CLEANUP+=("$dbname:$slot")
 142 }
 143
 144 #
 145 # Test cases
 146 #
 147
 148 test_basic_restore_from_base_only() {
 149     ((TESTS_RUN++))
 150     log_test "Basic restore from base backup only"
 151
 152     local source_db="${TEST_DB_PREFIX}_source_basic"
 153     local restore_db="${TEST_DB_PREFIX}_restore_basic"
 154     local slot="test_slot_basic"
 155     local backup_dir="$TEST_DIR/basic"
 156
 157     # Setup source database
 158     create_test_db "$source_db"
 159     create_table_with_pk "$source_db" "users"
 160     query_db "$source_db" "INSERT INTO users (name) VALUES ('Alice'), ('Bob'), ('Charlie');"
 161
 162     # Initialize backup
 163     mkdir -p "$backup_dir"
 164     initialize_backup "$source_db" "$backup_dir" "$slot"
 165
 166     # Restore to new database with --create
 167     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 168         DATABASES_TO_CLEANUP+=("$restore_db")
 169
 170         # Verify data
 171         local count
 172         count=$(query_db "$restore_db" "SELECT COUNT(*) FROM users;")
 173         if [[ "$count" -ne 3 ]]; then
 174             log_fail "Expected 3 rows, got $count"
 175             return 1
 176         fi
 177
 178         # Verify specific data (psql returns one value per line)
 179         local names
 180         names=$(query_db "$restore_db" "SELECT name FROM users ORDER BY id;")
 181         local expected_names=$'Alice\nBob\nCharlie'
 182         if [[ "$names" != "$expected_names" ]]; then
 183             log_fail "Data mismatch: expected '$expected_names', got '$names'"
 184             return 1
 185         fi
 186
 187         log_pass "Basic restore successful"
 188         return 0
 189     else
 190         log_fail "Restore command failed"
 191         return 1
 192     fi
 193 }
 194
 195 test_restore_with_incremental_backups() {
 196     ((TESTS_RUN++))
 197     log_test "Restore with base + incremental backups"
 198
 199     local source_db="${TEST_DB_PREFIX}_source_incr"
 200     local restore_db="${TEST_DB_PREFIX}_restore_incr"
 201     local slot="test_slot_incr"
 202     local backup_dir="$TEST_DIR/incremental"
 203
 204     # Setup source database
 205     create_test_db "$source_db"
 206     create_table_with_pk "$source_db" "orders"
 207     query_db "$source_db" "INSERT INTO orders (name) VALUES ('Order1'), ('Order2');"
 208
 209     # Initialize backup
 210     mkdir -p "$backup_dir"
 211     initialize_backup "$source_db" "$backup_dir" "$slot"
 212
 213     # Start collecting incremental changes
 214     local incr_file="$backup_dir/incremental.sql"
 215     "$PG_SCRIBE" --start -d "$source_db" -f "$incr_file" -S "$slot" -U "$PGUSER" &>/dev/null &
 216     local pg_scribe_pid=$!
 217
 218     # Wait for streaming to start
 219     sleep 2
 220
 221     # Make some changes
 222     query_db "$source_db" "INSERT INTO orders (name) VALUES ('Order3');"
 223     query_db "$source_db" "UPDATE orders SET name = 'Order1_Updated' WHERE name = 'Order1';"
 224     query_db "$source_db" "DELETE FROM orders WHERE name = 'Order2';"
 225
 226     # Wait for changes to be captured
 227     sleep 2
 228
 229     # Stop streaming
 230     kill "$pg_scribe_pid" 2>/dev/null || true
 231     wait "$pg_scribe_pid" 2>/dev/null || true
 232
 233     # Verify incremental file exists and has content
 234     if [[ ! -s "$incr_file" ]]; then
 235         log_fail "Incremental backup file is empty or missing"
 236         return 1
 237     fi
 238
 239     # Restore to new database
 240     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 241         DATABASES_TO_CLEANUP+=("$restore_db")
 242
 243         # Verify final state (should have Order1_Updated and Order3, not Order2)
 244         local count
 245         count=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders;")
 246         if [[ "$count" -ne 2 ]]; then
 247             log_fail "Expected 2 rows after incremental restore, got $count"
 248             return 1
 249         fi
 250
 251         # Verify Order1 was updated
 252         local order1_name
 253         order1_name=$(query_db "$restore_db" "SELECT name FROM orders WHERE id = 1;")
 254         if [[ "$order1_name" != "Order1_Updated" ]]; then
 255             log_fail "UPDATE not applied: expected 'Order1_Updated', got '$order1_name'"
 256             return 1
 257         fi
 258
 259         # Verify Order3 exists
 260         local order3_exists
 261         order3_exists=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders WHERE name = 'Order3';")
 262         if [[ "$order3_exists" -ne 1 ]]; then
 263             log_fail "INSERT not applied: Order3 not found"
 264             return 1
 265         fi
 266
 267         # Verify Order2 was deleted
 268         local order2_exists
 269         order2_exists=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders WHERE name = 'Order2';")
 270         if [[ "$order2_exists" -ne 0 ]]; then
 271             log_fail "DELETE not applied: Order2 still exists"
 272             return 1
 273         fi
 274
 275         log_pass "Restore with incrementals successful"
 276         return 0
 277     else
 278         log_fail "Restore command failed"
 279         return 1
 280     fi
 281 }
 282
 283 test_restore_with_ddl_changes() {
 284     ((TESTS_RUN++))
 285     log_test "Restore with DDL changes in incremental"
 286
 287     local source_db="${TEST_DB_PREFIX}_source_ddl"
 288     local restore_db="${TEST_DB_PREFIX}_restore_ddl"
 289     local slot="test_slot_ddl"
 290     local backup_dir="$TEST_DIR/ddl"
 291
 292     # Setup source database
 293     create_test_db "$source_db"
 294     create_table_with_pk "$source_db" "products"
 295     query_db "$source_db" "INSERT INTO products (name) VALUES ('Widget');"
 296
 297     # Initialize backup
 298     mkdir -p "$backup_dir"
 299     initialize_backup "$source_db" "$backup_dir" "$slot"
 300
 301     # Start collecting incremental changes
 302     local incr_file="$backup_dir/incremental.sql"
 303     "$PG_SCRIBE" --start -d "$source_db" -f "$incr_file" -S "$slot" -U "$PGUSER" &>/dev/null &
 304     local pg_scribe_pid=$!
 305
 306     # Wait for streaming to start
 307     sleep 2
 308
 309     # Add a column (DDL change)
 310     query_db "$source_db" "ALTER TABLE products ADD COLUMN price NUMERIC(10,2);"
 311
 312     # Insert data using new column
 313     query_db "$source_db" "INSERT INTO products (name, price) VALUES ('Gadget', 19.99);"
 314
 315     # Wait for changes to be captured
 316     sleep 2
 317
 318     # Stop streaming
 319     kill "$pg_scribe_pid" 2>/dev/null || true
 320     wait "$pg_scribe_pid" 2>/dev/null || true
 321
 322     # Restore to new database
 323     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 324         DATABASES_TO_CLEANUP+=("$restore_db")
 325
 326         # Verify new column exists
 327         local has_price_column
 328         has_price_column=$(query_db "$restore_db" "
 329             SELECT COUNT(*) FROM information_schema.columns
 330             WHERE table_name = 'products' AND column_name = 'price';
 331         ")
 332         if [[ "$has_price_column" -ne 1 ]]; then
 333             log_fail "DDL not applied: price column not found"
 334             return 1
 335         fi
 336
 337         # Verify data with new column
 338         local gadget_price
 339         gadget_price=$(query_db "$restore_db" "SELECT price FROM products WHERE name = 'Gadget';")
 340         if [[ "$gadget_price" != "19.99" ]]; then
 341             log_fail "Data with new column not correct: expected '19.99', got '$gadget_price'"
 342             return 1
 343         fi
 344
 345         log_pass "Restore with DDL changes successful"
 346         return 0
 347     else
 348         log_fail "Restore command failed"
 349         return 1
 350     fi
 351 }
 352
 353 test_restore_sequence_synchronization() {
 354     ((TESTS_RUN++))
 355     log_test "Restore with sequence synchronization"
 356
 357     local source_db="${TEST_DB_PREFIX}_source_seq"
 358     local restore_db="${TEST_DB_PREFIX}_restore_seq"
 359     local slot="test_slot_seq"
 360     local backup_dir="$TEST_DIR/sequence"
 361
 362     # Setup source database
 363     create_test_db "$source_db"
 364     create_table_with_pk "$source_db" "items"
 365     query_db "$source_db" "INSERT INTO items (name) VALUES ('Item1'), ('Item2'), ('Item3');"
 366
 367     # Initialize backup
 368     mkdir -p "$backup_dir"
 369     initialize_backup "$source_db" "$backup_dir" "$slot"
 370
 371     # Add more data to advance sequence
 372     query_db "$source_db" "INSERT INTO items (name) VALUES ('Item4'), ('Item5');"
 373
 374     # Collect incremental
 375     local incr_file="$backup_dir/incremental.sql"
 376     "$PG_SCRIBE" --start -d "$source_db" -f "$incr_file" -S "$slot" -U "$PGUSER" &>/dev/null &
 377     local pg_scribe_pid=$!
 378     sleep 2
 379     kill "$pg_scribe_pid" 2>/dev/null || true
 380     wait "$pg_scribe_pid" 2>/dev/null || true
 381
 382     # Restore to new database
 383     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 384         DATABASES_TO_CLEANUP+=("$restore_db")
 385
 386         # Get current sequence value
 387         local seq_val
 388         seq_val=$(query_db "$restore_db" "SELECT last_value FROM items_id_seq;")
 389
 390         # Should be at least 5 (we inserted 5 items total)
 391         if [[ "$seq_val" -lt 5 ]]; then
 392             log_fail "Sequence not synchronized: expected >= 5, got $seq_val"
 393             return 1
 394         fi
 395
 396         # Try inserting new row - should get ID 6
 397         query_db "$restore_db" "INSERT INTO items (name) VALUES ('Item6');"
 398         local new_id
 399         new_id=$(query_db "$restore_db" "SELECT id FROM items WHERE name = 'Item6';")
 400         if [[ "$new_id" -ne 6 ]]; then
 401             log_fail "Next sequence value incorrect: expected 6, got $new_id"
 402             return 1
 403         fi
 404
 405         log_pass "Sequence synchronization successful"
 406         return 0
 407     else
 408         log_fail "Restore command failed"
 409         return 1
 410     fi
 411 }
 412
 413 test_restore_no_sync_sequences() {
 414     ((TESTS_RUN++))
 415     log_test "Restore with --no-sync-sequences flag"
 416
 417     local source_db="${TEST_DB_PREFIX}_source_noseq"
 418     local restore_db="${TEST_DB_PREFIX}_restore_noseq"
 419     local slot="test_slot_noseq"
 420     local backup_dir="$TEST_DIR/noseq"
 421
 422     # Setup source database
 423     create_test_db "$source_db"
 424     create_table_with_pk "$source_db" "records"
 425     query_db "$source_db" "INSERT INTO records (name) VALUES ('Rec1'), ('Rec2');"
 426
 427     # Initialize backup
 428     mkdir -p "$backup_dir"
 429     initialize_backup "$source_db" "$backup_dir" "$slot"
 430
 431     # Restore with --no-sync-sequences
 432     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" --no-sync-sequences &>/dev/null; then
 433         DATABASES_TO_CLEANUP+=("$restore_db")
 434
 435         # Sequence should be at the value from pg_dump (2)
 436         local seq_val
 437         seq_val=$(query_db "$restore_db" "SELECT last_value FROM records_id_seq;")
 438
 439         # With --no-sync-sequences, sequence value is from pg_dump
 440         # It should match what's in the backup
 441         if [[ "$seq_val" -lt 1 ]]; then
 442             log_fail "Sequence value unexpectedly low: $seq_val"
 443             return 1
 444         fi
 445
 446         log_pass "Restore with --no-sync-sequences successful"
 447         return 0
 448     else
 449         log_fail "Restore command failed"
 450         return 1
 451     fi
 452 }
 453
 454 test_restore_to_existing_database() {
 455     ((TESTS_RUN++))
 456     log_test "Restore to existing database (without --create)"
 457
 458     local source_db="${TEST_DB_PREFIX}_source_exist"
 459     local restore_db="${TEST_DB_PREFIX}_restore_exist"
 460     local slot="test_slot_exist"
 461     local backup_dir="$TEST_DIR/existing"
 462
 463     # Setup source database
 464     create_test_db "$source_db"
 465     create_table_with_pk "$source_db" "data"
 466     query_db "$source_db" "INSERT INTO data (name) VALUES ('Test1');"
 467
 468     # Initialize backup
 469     mkdir -p "$backup_dir"
 470     initialize_backup "$source_db" "$backup_dir" "$slot"
 471
 472     # Pre-create target database
 473     create_test_db "$restore_db"
 474
 475     # Restore without --create
 476     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -U "$PGUSER" &>/dev/null; then
 477         # Verify data
 478         local count
 479         count=$(query_db "$restore_db" "SELECT COUNT(*) FROM data;")
 480         if [[ "$count" -ne 1 ]]; then
 481             log_fail "Expected 1 row, got $count"
 482             return 1
 483         fi
 484
 485         log_pass "Restore to existing database successful"
 486         return 0
 487     else
 488         log_fail "Restore command failed"
 489         return 1
 490     fi
 491 }
 492
 493 test_restore_fails_if_db_exists_with_create() {
 494     ((TESTS_RUN++))
 495     log_test "Restore fails if database exists with --create"
 496
 497     local source_db="${TEST_DB_PREFIX}_source_exists"
 498     local restore_db="${TEST_DB_PREFIX}_restore_exists"
 499     local slot="test_slot_exists"
 500     local backup_dir="$TEST_DIR/exists"
 501
 502     # Setup source database
 503     create_test_db "$source_db"
 504     create_table_with_pk "$source_db" "test"
 505     query_db "$source_db" "INSERT INTO test (name) VALUES ('Test');"
 506
 507     # Initialize backup
 508     mkdir -p "$backup_dir"
 509     initialize_backup "$source_db" "$backup_dir" "$slot"
 510
 511     # Pre-create target database
 512     create_test_db "$restore_db"
 513
 514     # Restore with --create should fail
 515     local exit_code=0
 516     "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null || exit_code=$?
 517
 518     if [[ $exit_code -eq 4 ]]; then
 519         log_pass "Correctly failed when database exists with --create"
 520         return 0
 521     else
 522         log_fail "Expected exit code 4, got $exit_code"
 523         return 1
 524     fi
 525 }
 526
 527 test_restore_missing_backup_directory() {
 528     ((TESTS_RUN++))
 529     log_test "Restore fails with missing backup directory"
 530
 531     local restore_db="${TEST_DB_PREFIX}_restore_missing"
 532     local backup_dir="$TEST_DIR/nonexistent"
 533
 534     # Try to restore from non-existent directory
 535     local exit_code=0
 536     "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -U "$PGUSER" &>/dev/null || exit_code=$?
 537
 538     if [[ $exit_code -eq 4 ]]; then
 539         log_pass "Correctly failed with missing backup directory"
 540         return 0
 541     else
 542         log_fail "Expected exit code 4, got $exit_code"
 543         return 1
 544     fi
 545 }
 546
 547 test_restore_specific_base_backup() {
 548     ((TESTS_RUN++))
 549     log_test "Restore specific base backup with --base-backup"
 550
 551     local source_db="${TEST_DB_PREFIX}_source_specific"
 552     local restore_db="${TEST_DB_PREFIX}_restore_specific"
 553     local slot="test_slot_specific"
 554     local backup_dir="$TEST_DIR/specific"
 555
 556     # Setup source database
 557     create_test_db "$source_db"
 558     create_table_with_pk "$source_db" "entries"
 559     query_db "$source_db" "INSERT INTO entries (name) VALUES ('Entry1');"
 560
 561     # Initialize backup (creates first base backup)
 562     mkdir -p "$backup_dir"
 563     initialize_backup "$source_db" "$backup_dir" "$slot"
 564
 565     # Find the base backup file
 566     local base_file
 567     base_file=$(find "$backup_dir" -maxdepth 1 -name 'base-*.sql' -print -quit)
 568
 569     if [[ -z "$base_file" ]]; then
 570         log_fail "Base backup file not found"
 571         return 1
 572     fi
 573
 574     # Restore using specific base backup
 575     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" --base-backup="$base_file" -C -U "$PGUSER" &>/dev/null; then
 576         DATABASES_TO_CLEANUP+=("$restore_db")
 577
 578         # Verify data
 579         local count
 580         count=$(query_db "$restore_db" "SELECT COUNT(*) FROM entries;")
 581         if [[ "$count" -ne 1 ]]; then
 582             log_fail "Expected 1 row, got $count"
 583             return 1
 584         fi
 585
 586         log_pass "Restore with specific base backup successful"
 587         return 0
 588     else
 589         log_fail "Restore command failed"
 590         return 1
 591     fi
 592 }
 593
 594 test_restore_multiple_tables() {
 595     ((TESTS_RUN++))
 596     log_test "Restore multiple tables with relationships"
 597
 598     local source_db="${TEST_DB_PREFIX}_source_multi"
 599     local restore_db="${TEST_DB_PREFIX}_restore_multi"
 600     local slot="test_slot_multi"
 601     local backup_dir="$TEST_DIR/multi"
 602
 603     # Setup source database with multiple related tables
 604     create_test_db "$source_db"
 605     query_db "$source_db" "
 606         CREATE TABLE customers (
 607             id SERIAL PRIMARY KEY,
 608             name TEXT NOT NULL
 609         );
 610         CREATE TABLE orders (
 611             id SERIAL PRIMARY KEY,
 612             customer_id INTEGER REFERENCES customers(id),
 613             product TEXT NOT NULL
 614         );
 615     "
 616     query_db "$source_db" "INSERT INTO customers (name) VALUES ('Alice'), ('Bob');"
 617     query_db "$source_db" "INSERT INTO orders (customer_id, product) VALUES (1, 'Widget'), (2, 'Gadget');"
 618
 619     # Initialize backup
 620     mkdir -p "$backup_dir"
 621     initialize_backup "$source_db" "$backup_dir" "$slot"
 622
 623     # Restore
 624     if "$PG_SCRIBE" --restore -d "$restore_db" -f "$backup_dir" -C -U "$PGUSER" &>/dev/null; then
 625         DATABASES_TO_CLEANUP+=("$restore_db")
 626
 627         # Verify customers
 628         local customer_count
 629         customer_count=$(query_db "$restore_db" "SELECT COUNT(*) FROM customers;")
 630         if [[ "$customer_count" -ne 2 ]]; then
 631             log_fail "Expected 2 customers, got $customer_count"
 632             return 1
 633         fi
 634
 635         # Verify orders
 636         local order_count
 637         order_count=$(query_db "$restore_db" "SELECT COUNT(*) FROM orders;")
 638         if [[ "$order_count" -ne 2 ]]; then
 639             log_fail "Expected 2 orders, got $order_count"
 640             return 1
 641         fi
 642
 643         # Verify foreign key relationship
 644         local alice_orders
 645         alice_orders=$(query_db "$restore_db" "
 646             SELECT o.product FROM orders o
 647             JOIN customers c ON c.id = o.customer_id
 648             WHERE c.name = 'Alice';
 649         ")
 650         if [[ "$alice_orders" != "Widget" ]]; then
 651             log_fail "Foreign key relationship not preserved"
 652             return 1
 653         fi
 654
 655         log_pass "Multiple table restore successful"
 656         return 0
 657     else
 658         log_fail "Restore command failed"
 659         return 1
 660     fi
 661 }
 662
 663 #
 664 # Cleanup
 665 #
 666
 667 # shellcheck disable=SC2317  # Function called via trap handler
 668 cleanup() {
 669     log_info "Cleaning up test resources..."
 670
 671     # Stop any background pg_scribe processes
 672     pkill -f "pg_scribe.*--start" 2>/dev/null || true
 673     pkill -f "pg_recvlogical" 2>/dev/null || true
 674     sleep 1
 675
 676     # Drop replication slots
 677     for slot_info in "${SLOTS_TO_CLEANUP[@]}"; do
 678         IFS=':' read -r dbname slot <<< "$slot_info"
 679         drop_replication_slot "$dbname" "$slot" 2>/dev/null || true
 680     done
 681
 682     # Drop databases
 683     for dbname in "${DATABASES_TO_CLEANUP[@]}"; do
 684         drop_test_db "$dbname"
 685     done
 686
 687     # Remove test directory
 688     if [[ -d "$TEST_DIR" ]]; then
 689         rm -rf "$TEST_DIR"
 690     fi
 691
 692     log_info "Cleanup complete"
 693 }
 694
 695 #
 696 # Main test runner
 697 #
 698
 699 main() {
 700     echo "========================================"
 701     echo "pg_scribe --restore Test Suite"
 702     echo "========================================"
 703     echo ""
 704
 705     # Verify pg_scribe exists
 706     if [[ ! -x "$PG_SCRIBE" ]]; then
 707         echo "ERROR: pg_scribe not found or not executable: $PG_SCRIBE"
 708         exit 1
 709     fi
 710
 711     # Verify PostgreSQL is running
 712     if ! psql -U "$PGUSER" -d postgres -c "SELECT 1;" &>/dev/null; then
 713         echo "ERROR: Cannot connect to PostgreSQL"
 714         exit 1
 715     fi
 716
 717     # Verify wal_level is logical
 718     local wal_level
 719     wal_level=$(psql -U "$PGUSER" -d postgres -tAq -c "SHOW wal_level;")
 720     if [[ "$wal_level" != "logical" ]]; then
 721         echo "ERROR: wal_level must be 'logical', currently: $wal_level"
 722         echo "Update ~/.pgenv/pgsql/data/postgresql.conf and restart PostgreSQL"
 723         exit 1
 724     fi
 725
 726     # Verify wal2sql extension is available
 727     if ! psql -U "$PGUSER" -d postgres -c "CREATE EXTENSION IF NOT EXISTS wal2sql;" &>/dev/null; then
 728         echo "ERROR: wal2sql extension not available"
 729         echo "Build and install: cd wal2sql && make && make install"
 730         exit 1
 731     fi
 732
 733     # Create test directory
 734     mkdir -p "$TEST_DIR"
 735
 736     # Set up cleanup trap
 737     trap cleanup EXIT INT TERM
 738
 739     echo "Running tests..."
 740     echo ""
 741
 742     # Run all tests (use || true to prevent set -e from exiting)
 743     test_basic_restore_from_base_only || true
 744     test_restore_with_incremental_backups || true
 745     test_restore_with_ddl_changes || true
 746     test_restore_sequence_synchronization || true
 747     test_restore_no_sync_sequences || true
 748     test_restore_to_existing_database || true
 749     test_restore_fails_if_db_exists_with_create || true
 750     test_restore_missing_backup_directory || true
 751     test_restore_specific_base_backup || true
 752     test_restore_multiple_tables || true
 753
 754     # Summary
 755     echo ""
 756     echo "========================================"
 757     echo "Test Results"
 758     echo "========================================"
 759     echo "Tests run:    $TESTS_RUN"
 760     echo -e "Tests passed: ${GREEN}$TESTS_PASSED${NC}"
 761     echo -e "Tests failed: ${RED}$TESTS_FAILED${NC}"
 762     echo ""
 763
 764     if [[ $TESTS_FAILED -eq 0 ]]; then
 765         echo -e "${GREEN}All tests passed!${NC}"
 766         exit 0
 767     else
 768         echo -e "${RED}Some tests failed!${NC}"
 769         exit 1
 770     fi
 771 }
 772
 773 main "$@"