Function TestPartitionRing_ShuffleShard

ring/partition_ring_test.go:174–219

(t *testing.T)

}

func TestPartitionRing_ShuffleShard(t *testing.T) {
	t.Run("should honor the shard size", func(t *testing.T) {
		const numActivePartitions = 5

		opts := DefaultPartitionRingOptions()
		// Set a size so we can assert the options are preserved on update.
		opts.ShuffleShardCacheSize = 1
		ring := createPartitionRingWithPartitions(opts, numActivePartitions, 0, 0)

		// Request a shard size up to the number of existing partitions.
		for shardSize := 1; shardSize <= numActivePartitions; shardSize++ {
			subring, err := ring.ShuffleShard("tenant-id", shardSize)
			require.NoError(t, err)
			assert.Equal(t, shardSize, subring.PartitionsCount())
			assert.Equal(t, subring.PartitionsCount(), ring.ShuffleShardSize(shardSize))
		}

		// Request a shard size greater than the number of existing partitions.
		subring, err := ring.ShuffleShard("tenant-id", numActivePartitions+1)
		require.NoError(t, err)
		assert.Equal(t, numActivePartitions, subring.PartitionsCount())
		assert.Equal(t, subring.PartitionsCount(), ring.ShuffleShardSize(numActivePartitions+1))
		assert.Equal(t, opts, subring.opts)
	})

	t.Run("should never return INACTIVE or PENDING partitions", func(t *testing.T) {
		const (
			numActivePartitions   = 5
			numInactivePartitions = 5
			numPendingPartitions  = 5
		)

		ring := createPartitionRingWithPartitions(DefaultPartitionRingOptions(), numActivePartitions, numInactivePartitions, numPendingPartitions)

		// We test negative values of shardSize as well as sizes above current number of partition count.
		for shardSize := -5; shardSize <= ring.PartitionsCount()+5; shardSize++ {
			subring, err := ring.ShuffleShard("tenant-id", shardSize)
			require.NoError(t, err)

			for _, partition := range subring.Partitions() {
				assert.Equal(t, PartitionActive, partition.State)
			}
			assert.Equal(t, subring.PartitionsCount(), ring.ShuffleShardSize(shardSize))
		}
	})
}

// This test asserts on shard stability across multiple invocations and given the same input ring.
func TestPartitionRing_ShuffleShard_Stability(t *testing.T) {