kythe.io@v0.0.68-0.20240422202219-7225dbc01741/kythe/go/util/disksort/disksort_test.go (about)

     1  /*
     2   * Copyright 2015 The Kythe Authors. All rights reserved.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *   http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package disksort
    18  
    19  import (
    20  	"fmt"
    21  	"math/rand"
    22  	"strconv"
    23  	"testing"
    24  )
    25  
    26  type numLesser struct{}
    27  
    28  // Less implements the Lesser interface.
    29  func (numLesser) Less(a, b any) bool { return a.(int) < b.(int) }
    30  
    31  type numMarshaler struct{}
    32  
    33  // Marshal implements part of the Marshaler interface.
    34  func (numMarshaler) Marshal(x any) ([]byte, error) {
    35  	return []byte(strconv.Itoa(x.(int))), nil
    36  }
    37  
    38  // Unmarshal implements part of the Marshaler interface.
    39  func (numMarshaler) Unmarshal(rec []byte) (any, error) {
    40  	return strconv.Atoi(string(rec))
    41  }
    42  
    43  func TestMergeSorter(t *testing.T) {
    44  	// Sort 1M numbers in chunks of 750 (~1.3k shards)
    45  	const n = 1000000
    46  	const max = 750
    47  
    48  	rand.Seed(120875)
    49  
    50  	sorter, err := NewMergeSorter(MergeOptions{
    51  		Lesser:      numLesser{},
    52  		Marshaler:   numMarshaler{},
    53  		MaxInMemory: max,
    54  	})
    55  	if err != nil {
    56  		t.Fatalf("error creating MergeSorter: %v", err)
    57  	}
    58  
    59  	nums := make([]int, n)
    60  	for i := 0; i < n; i++ {
    61  		nums[i] = i
    62  	}
    63  
    64  	// Randomize order for Add
    65  	for i := 0; i < n; i++ {
    66  		o := rand.Int() % n
    67  		nums[i], nums[o] = nums[o], nums[i]
    68  	}
    69  
    70  	for _, n := range nums {
    71  		if err := sorter.Add(n); err != nil {
    72  			t.Fatalf("error adding %d to sorter: %v", n, err)
    73  		}
    74  	}
    75  
    76  	var expected int
    77  	if err := sorter.Read(func(i any) error {
    78  		x, ok := i.(int)
    79  		if !ok {
    80  			return fmt.Errorf("expected int; found %T", i)
    81  		} else if expected != x {
    82  			return fmt.Errorf("expected %d; found %d", expected, x)
    83  		}
    84  		expected++
    85  		return nil
    86  	}); err != nil {
    87  		t.Fatalf("read error: %v", err)
    88  	}
    89  
    90  	if expected != n {
    91  		t.Fatalf("Expected %d total; found %d", n, expected)
    92  	}
    93  }