"""Unit tests for db_utils.""" from dataclasses import dataclass import pytest from src.utils.db_utils import build_batches, calculate_max_batch_size class TestCalculateMaxBatchSize: """Tests for calculate_max_batch_size.""" def test_basic_calculation(self): """Test basic batch size calculation.""" # 10_000 bytes, 10 bytes/entry, 100 base, 100% margin # available = 10_000 - 100 = 9_900 # (9_900 + 1) / (10 + 1) = 900 result = calculate_max_batch_size( max_query_bytes=10_000, bytes_per_entry=10, base_query_bytes=100, safety_margin_pct=1.0, ) assert result == 900 def test_safety_margin(self): """Test safety margin reduces available bytes.""" # 10_000 bytes * 0.5 margin = 5_000 effective, minus 100 base = 4_900 # (4_900 + 1) / (10 + 1) = 445 result = calculate_max_batch_size( max_query_bytes=10_000, bytes_per_entry=10, base_query_bytes=100, safety_margin_pct=0.5, ) assert result == 445 def test_default_safety_margin(self): """Test default safety margin is 0.9.""" # 10_000 * 0.9 = 9_000, minus 100 = 8_900 # (8_900 + 1) / (10 + 1) = 809 result = calculate_max_batch_size( max_query_bytes=10_000, bytes_per_entry=10, base_query_bytes=100, ) assert result == 809 def test_returns_zero_when_nothing_fits(self): """Test returns 0 when not even one entry fits.""" result = calculate_max_batch_size( max_query_bytes=50, bytes_per_entry=100, base_query_bytes=0, safety_margin_pct=1.0, ) assert result == 0 def test_large_base_query(self): """Test large base query leaves little room.""" # 1_000 * 1.0 = 1_000, minus 990 = 10 available # (10 + 1) / (20 + 1) = 0 result = calculate_max_batch_size( max_query_bytes=1_000, bytes_per_entry=20, base_query_bytes=990, safety_margin_pct=1.0, ) assert result == 0 def test_realistic_mysql_values(self): """Test with realistic MySQL max_allowed_packet (16MB).""" # 16MB * 0.9 = 14.4MB, minus 1000 base ≈ 14.4MB # ~14.4M / 21 ≈ 685_714 result = calculate_max_batch_size( max_query_bytes=16_777_216, bytes_per_entry=20, base_query_bytes=1000, safety_margin_pct=0.9, ) assert result == 718_975 def test_invalid_max_query_bytes(self): """Test raises ValueError for non-positive max_query_bytes.""" with pytest.raises(ValueError, match='max_query_bytes must be positive'): calculate_max_batch_size( max_query_bytes=0, bytes_per_entry=10, base_query_bytes=0 ) def test_invalid_bytes_per_entry(self): """Test raises ValueError for non-positive bytes_per_entry.""" with pytest.raises(ValueError, match='bytes_per_entry must be positive'): calculate_max_batch_size( max_query_bytes=1000, bytes_per_entry=0, base_query_bytes=0 ) def test_invalid_base_query_bytes(self): """Test raises ValueError for negative base_query_bytes.""" with pytest.raises(ValueError, match='base_query_bytes must be nonnegative'): calculate_max_batch_size( max_query_bytes=1000, bytes_per_entry=10, base_query_bytes=-1 ) def test_invalid_safety_margin_zero(self): """Test raises ValueError for zero safety_margin_pct.""" with pytest.raises(ValueError, match='safety_margin_pct must be within'): calculate_max_batch_size( max_query_bytes=1000, bytes_per_entry=10, base_query_bytes=0, safety_margin_pct=0.0, ) def test_invalid_safety_margin_above_one(self): """Test raises ValueError for safety_margin_pct > 1.""" with pytest.raises(ValueError, match='safety_margin_pct must be within'): calculate_max_batch_size( max_query_bytes=1000, bytes_per_entry=10, base_query_bytes=0, safety_margin_pct=1.1, ) @dataclass class _Item: """Simple item for testing build_batches.""" id: int weight: int class TestBuildBatches: """Tests for build_batches.""" def test_single_item_fits_in_batch(self): """Test single item packs into one batch.""" items = [_Item(id=1, weight=100)] batches = build_batches( items, 1000, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(1, 100)]] def test_multiple_items_single_batch(self): """Test multiple items fit in a single batch (sorted by size desc).""" items = [ _Item(id=1, weight=300), _Item(id=2, weight=400), _Item(id=3, weight=200), ] batches = build_batches( items, 1000, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(2, 400), (1, 300), (3, 200)]] def test_multiple_batches(self): """Test items split across batches; FFD packs 400 into first batch with 600.""" items = [ _Item(id=1, weight=600), _Item(id=2, weight=600), _Item(id=3, weight=400), ] batches = build_batches( items, 1000, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(1, 600), (3, 400)], [(2, 600)]] def test_oversized_item(self): """Test oversized item splits into full portion + remainder.""" items = [_Item(id=1, weight=1500)] batches = build_batches( items, 1000, key=lambda x: x.id, size=lambda x: x.weight, ) # 1500 / 1000 = 1 full portion (1000), remainder 500 assert batches == [[(1, 1000)], [(1, 500)]] def test_mix_regular_and_oversized(self): """Test oversized remainder is FFD-packed with regular items. Sorted: (2, 1500), (3, 400), (1, 300). Item 2 is oversized: full portion (2, 1000), remainder (2, 500) enters FFD. Then (3, 400) doesn't fit with remainder (500 remaining), new batch. Then (1, 300) doesn't fit with remainder or with (3), new batch. """ items = [ _Item(id=1, weight=300), _Item(id=2, weight=1500), _Item(id=3, weight=400), ] batches = build_batches( items, 1000, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(2, 1000)], [(2, 500), (3, 400)], [(1, 300)]] def test_empty_input(self): """Test empty input returns empty results.""" batches = build_batches( [], 1000, key=lambda x: x, size=lambda x: x, ) assert batches == [] def test_exact_batch_boundary(self): """Test items exactly filling max_batch_size start a new batch on overflow.""" items = [ _Item(id=1, weight=500), _Item(id=2, weight=500), _Item(id=3, weight=100), ] batches = build_batches( items, 1000, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(1, 500), (2, 500)], [(3, 100)]] def test_item_size_equals_max_batch_size_fits(self): """Test item with size exactly equal to max_batch_size fits in a batch.""" items = [_Item(id=1, weight=1000)] batches = build_batches( items, 1000, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(1, 1000)]] def test_ffd_fills_gaps(self): """Test FFD fills gaps that greedy sequential would miss. Greedy sequential would produce 3 batches: [[1], [2, 3], [4]]. FFD sorts [600, 600, 400, 350] and packs into 2 batches. """ items = [ _Item(id=1, weight=600), _Item(id=2, weight=600), _Item(id=3, weight=400), _Item(id=4, weight=350), ] batches = build_batches( items, 1000, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(1, 600), (3, 400)], [(2, 600), (4, 350)]] def test_oversized_remainder_packed_with_regular(self): """Test oversized remainder is packed with a small regular item. Item (1, 1050) with max=100: full portion (1, 1000), remainder (1, 50). Item (2, 40) fits in the remainder batch (50 remaining after remainder). Saves one round-trip vs processing remainder alone. """ items = [ _Item(id=1, weight=1050), _Item(id=2, weight=40), ] batches = build_batches( items, 100, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(1, 1000)], [(1, 50), (2, 40)]] def test_oversized_exact_multiple(self): """Test oversized item that is an exact multiple has no remainder.""" items = [_Item(id=1, weight=200)] batches = build_batches( items, 100, key=lambda x: x.id, size=lambda x: x.weight, ) assert batches == [[(1, 200)]]