@@ -323,4 +323,233 @@ mod tests {
323323 assert_eq ! ( PartitionStrategy :: from_str( "adaptive" ) , Some ( PartitionStrategy :: Adaptive ) ) ;
324324 assert_eq ! ( PartitionStrategy :: from_str( "invalid" ) , None ) ;
325325 }
326+
327+ // -----------------------------------------------------------------------
328+ // Additional partition tests
329+ // -----------------------------------------------------------------------
330+
331+ /// Partition 1 item across many locales — only the first locale gets the
332+ /// item, all others get zero-length slices. The partition must still be
333+ /// complete and non-overlapping.
334+ #[ test]
335+ fn per_item_partition_1_item_many_locales ( ) {
336+ let p = Partition :: per_item ( 1 , 16 ) ;
337+ assert ! ( p. verify( ) , "1 item across 16 locales should be valid" ) ;
338+ assert_eq ! ( p. slices. len( ) , 16 ) ;
339+ // Exactly one locale gets the item
340+ assert_eq ! ( p. slices[ 0 ] . count, 1 , "First locale should get the 1 item" ) ;
341+ for s in & p. slices [ 1 ..] {
342+ assert_eq ! ( s. count, 0 , "Remaining locales should have 0 items" ) ;
343+ }
344+ }
345+
346+ /// Partition a prime number of items — ensures the remainder distribution
347+ /// is correct (first N%K locales get one extra).
348+ #[ test]
349+ fn per_item_partition_prime_items ( ) {
350+ // 97 items across 8 locales: 97/8 = 12 base, 97%8 = 1 remainder
351+ let p = Partition :: per_item ( 97 , 8 ) ;
352+ assert ! ( p. verify( ) , "97 items across 8 locales should be valid" ) ;
353+ assert_eq ! ( p. slices. len( ) , 8 ) ;
354+ assert_eq ! ( p. slices[ 0 ] . count, 13 , "First locale gets 12+1 = 13" ) ;
355+ for s in & p. slices [ 1 ..] {
356+ assert_eq ! ( s. count, 12 , "Remaining locales get 12 each" ) ;
357+ }
358+ // Verify total: 13 + 7*12 = 13 + 84 = 97
359+ let total: u64 = p. slices . iter ( ) . map ( |s| s. count ) . sum ( ) ;
360+ assert_eq ! ( total, 97 ) ;
361+ }
362+
363+ /// Another prime: 13 items across 5 locales (13/5 = 2 base, 3 remainder).
364+ #[ test]
365+ fn per_item_partition_prime_items_2 ( ) {
366+ let p = Partition :: per_item ( 13 , 5 ) ;
367+ assert ! ( p. verify( ) ) ;
368+ // First 3 locales get 3, last 2 get 2
369+ assert_eq ! ( p. slices[ 0 ] . count, 3 ) ;
370+ assert_eq ! ( p. slices[ 1 ] . count, 3 ) ;
371+ assert_eq ! ( p. slices[ 2 ] . count, 3 ) ;
372+ assert_eq ! ( p. slices[ 3 ] . count, 2 ) ;
373+ assert_eq ! ( p. slices[ 4 ] . count, 2 ) ;
374+ }
375+
376+ /// Chunked partition with 1 item — single chunk, single locale gets it.
377+ #[ test]
378+ fn chunked_partition_1_item ( ) {
379+ let p = Partition :: chunked ( 1 , 4 , 10 ) ;
380+ assert ! ( p. verify( ) , "Chunked partition of 1 item should be valid" ) ;
381+ let non_empty: Vec < _ > = p. slices . iter ( ) . filter ( |s| s. count > 0 ) . collect ( ) ;
382+ assert_eq ! ( non_empty. len( ) , 1 , "Only one locale should have the item" ) ;
383+ assert_eq ! ( non_empty[ 0 ] . count, 1 ) ;
384+ }
385+
386+ /// Chunked partition where grain_size > total_items.
387+ #[ test]
388+ fn chunked_partition_large_grain ( ) {
389+ let p = Partition :: chunked ( 5 , 4 , 100 ) ;
390+ assert ! ( p. verify( ) , "Chunked partition with large grain should be valid" ) ;
391+ // All items in one chunk on one locale
392+ let total: u64 = p. slices . iter ( ) . map ( |s| s. count ) . sum ( ) ;
393+ assert_eq ! ( total, 5 ) ;
394+ }
395+
396+ // -----------------------------------------------------------------------
397+ // Additional GatherResult tests
398+ // -----------------------------------------------------------------------
399+
400+ /// GatherResult with 0 total results and empty locale_counts.
401+ #[ test]
402+ fn gather_result_zero_results ( ) {
403+ let g = GatherResult {
404+ total_results : 0 ,
405+ locale_counts : vec ! [ ] ,
406+ strategy : GatherStrategy :: Merge ,
407+ } ;
408+ assert ! (
409+ g. verify_conservation( ) ,
410+ "0 results with empty locale_counts should be conserved"
411+ ) ;
412+ }
413+
414+ /// GatherResult with 0 total results but non-empty (all zero) locale_counts.
415+ #[ test]
416+ fn gather_result_zero_results_many_locales ( ) {
417+ let g = GatherResult {
418+ total_results : 0 ,
419+ locale_counts : vec ! [ 0 , 0 , 0 , 0 ] ,
420+ strategy : GatherStrategy :: Reduce ,
421+ } ;
422+ assert ! (
423+ g. verify_conservation( ) ,
424+ "0 results across 4 locales should be conserved"
425+ ) ;
426+ }
427+
428+ /// GatherResult conservation fails when total_results != sum(locale_counts).
429+ #[ test]
430+ fn gather_result_conservation_fails ( ) {
431+ let g = GatherResult {
432+ total_results : 100 ,
433+ locale_counts : vec ! [ 25 , 25 , 25 , 24 ] , // sum=99, not 100
434+ strategy : GatherStrategy :: Merge ,
435+ } ;
436+ assert ! (
437+ !g. verify_conservation( ) ,
438+ "Mismatched total should fail conservation check"
439+ ) ;
440+ }
441+
442+ /// GatherResult with a single locale holding all results.
443+ #[ test]
444+ fn gather_result_single_locale ( ) {
445+ let g = GatherResult {
446+ total_results : 42 ,
447+ locale_counts : vec ! [ 42 ] ,
448+ strategy : GatherStrategy :: First ,
449+ } ;
450+ assert ! ( g. verify_conservation( ) ) ;
451+ }
452+
453+ // -----------------------------------------------------------------------
454+ // Additional MemoryBudget tests
455+ // -----------------------------------------------------------------------
456+
457+ /// MemoryBudget with 0 max_item_bytes — everything should be zero except
458+ /// the metadata overhead.
459+ #[ test]
460+ fn memory_budget_zero_item_bytes ( ) {
461+ let budget = MemoryBudget :: calculate ( 100 , 4 , 0 ) ;
462+ assert_eq ! ( budget. input_bytes, 0 , "0 max_item_bytes means 0 input buffer" ) ;
463+ assert_eq ! ( budget. output_bytes, 0 , "0 max_item_bytes means 0 output buffer" ) ;
464+ assert ! ( budget. metadata_bytes > 0 , "Metadata bytes should still be non-zero" ) ;
465+ assert_eq ! (
466+ budget. total_bytes,
467+ budget. metadata_bytes,
468+ "Total should be metadata only"
469+ ) ;
470+ }
471+
472+ /// MemoryBudget with 1 locale — all items on one node.
473+ #[ test]
474+ fn memory_budget_single_locale ( ) {
475+ let budget = MemoryBudget :: calculate ( 100 , 1 , 1024 ) ;
476+ // items_per_locale = 100/1 + 1 = 101
477+ assert_eq ! ( budget. items_per_locale, 101 ) ;
478+ assert_eq ! ( budget. input_bytes, 101 * 1024 ) ;
479+ assert_eq ! ( budget. output_bytes, 101 * 1024 ) ;
480+ }
481+
482+ /// MemoryBudget with very large item size — check total_mb calculation.
483+ #[ test]
484+ fn memory_budget_large_items ( ) {
485+ // 10 items, 2 locales, 100MB per item
486+ let budget = MemoryBudget :: calculate ( 10 , 2 , 100 * 1_048_576 ) ;
487+ // items_per_locale = 10/2 + 1 = 6
488+ assert_eq ! ( budget. items_per_locale, 6 ) ;
489+ // input = 6 * 100MB = 600MB, output = 600MB, meta = 6*9 = 54
490+ assert_eq ! ( budget. input_bytes, 6 * 100 * 1_048_576 ) ;
491+ assert ! ( budget. total_mb >= 1200 , "Should be at least 1200MB" ) ;
492+ }
493+
494+ /// MemoryBudget with 1 item and many locales — items_per_locale should be 1.
495+ #[ test]
496+ fn memory_budget_1_item_many_locales ( ) {
497+ let budget = MemoryBudget :: calculate ( 1 , 100 , 4096 ) ;
498+ // items_per_locale = 1/100 + 1 = 0 + 1 = 1
499+ assert_eq ! ( budget. items_per_locale, 1 ) ;
500+ assert_eq ! ( budget. input_bytes, 4096 ) ;
501+ assert_eq ! ( budget. output_bytes, 4096 ) ;
502+ }
503+
504+ // -----------------------------------------------------------------------
505+ // Additional strategy parsing tests
506+ // -----------------------------------------------------------------------
507+
508+ /// All five partition strategy strings parse correctly.
509+ #[ test]
510+ fn partition_strategy_all_valid ( ) {
511+ let pairs = [
512+ ( "per-item" , PartitionStrategy :: PerItem ) ,
513+ ( "chunk" , PartitionStrategy :: Chunk ) ,
514+ ( "adaptive" , PartitionStrategy :: Adaptive ) ,
515+ ( "spatial" , PartitionStrategy :: Spatial ) ,
516+ ( "keyed" , PartitionStrategy :: Keyed ) ,
517+ ] ;
518+ for ( s, expected) in & pairs {
519+ assert_eq ! (
520+ PartitionStrategy :: from_str( s) ,
521+ Some ( * expected) ,
522+ "'{s}' should parse to {expected:?}"
523+ ) ;
524+ }
525+ }
526+
527+ /// All five gather strategy strings parse correctly.
528+ #[ test]
529+ fn gather_strategy_all_valid ( ) {
530+ let pairs = [
531+ ( "merge" , GatherStrategy :: Merge ) ,
532+ ( "reduce" , GatherStrategy :: Reduce ) ,
533+ ( "tree-reduce" , GatherStrategy :: TreeReduce ) ,
534+ ( "stream" , GatherStrategy :: Stream ) ,
535+ ( "first" , GatherStrategy :: First ) ,
536+ ] ;
537+ for ( s, expected) in & pairs {
538+ assert_eq ! (
539+ GatherStrategy :: from_str( s) ,
540+ Some ( * expected) ,
541+ "'{s}' should parse to {expected:?}"
542+ ) ;
543+ }
544+ }
545+
546+ /// Invalid strategy strings return None for both partition and gather.
547+ #[ test]
548+ fn strategy_parsing_rejects_invalid ( ) {
549+ let invalids = [ "" , "MERGE" , "per_item" , "Per-Item" , "tree_reduce" , " merge" ] ;
550+ for s in & invalids {
551+ assert_eq ! ( PartitionStrategy :: from_str( s) , None , "Partition should reject '{s}'" ) ;
552+ assert_eq ! ( GatherStrategy :: from_str( s) , None , "Gather should reject '{s}'" ) ;
553+ }
554+ }
326555}
0 commit comments