github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/diskqueue/diskqueue_test.go (about)

     1  package diskqueue
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"testing"
     7  
     8  	"gotest.tools/assert"
     9  )
    10  
    11  func TestQueue(t *testing.T) {
    12  	conf := Conf{Directory: "/tmp/dq", WriteMmap: true}
    13  	q, err := New(conf)
    14  	assert.Assert(t, err == nil)
    15  
    16  	testData := []byte("abcd")
    17  
    18  	// test Read
    19  	n := 1000
    20  	var offsets []int64
    21  	for i := 0; i < n; i++ {
    22  		offset, err := q.Put(testData)
    23  		offsets = append(offsets, offset)
    24  		assert.Assert(t, err == nil)
    25  
    26  		readData, err := q.Read(nil, offset)
    27  		// fmt.Println(string(readData))
    28  		assert.Assert(t, err == nil && bytes.Equal(readData, testData), "%v:%v", err, i)
    29  		assert.Assert(t, q.FileMeta(0).MsgCount == uint64(i+1))
    30  	}
    31  
    32  	assert.Assert(t, q.NumFiles() == 1 && q.FileMeta(0).MsgCount == uint64(n))
    33  
    34  	// test StreamOffsetRead
    35  	offsetCh := make(chan int64)
    36  	ch, err := q.StreamOffsetRead(offsetCh)
    37  	go func() {
    38  		for i := 0; i < n; i++ {
    39  			offsetCh <- offsets[i]
    40  		}
    41  	}()
    42  	for i := 0; i < n; i++ {
    43  		streamData, ok := <-ch
    44  		assert.Assert(t, bytes.Equal(streamData.Bytes, testData) && streamData.Offset == offsets[i], "i:%v ok:%v streamData.Offset:%v offsets[i]:%v", i, ok, streamData.Offset, offsets[i])
    45  	}
    46  	close(offsetCh)
    47  
    48  	// test StreamRead
    49  	ch, err = q.StreamRead(context.Background(), 0)
    50  	assert.Assert(t, err == nil)
    51  	for i := 0; i < n; i++ {
    52  		streamData := <-ch
    53  		assert.Assert(t, bytes.Equal(streamData.Bytes, testData))
    54  	}
    55  
    56  	n, err = q.GC()
    57  	assert.Assert(t, err == nil && n == 0)
    58  
    59  	err = q.Delete()
    60  	assert.Assert(t, err == nil)
    61  
    62  }
    63  
    64  func TestFixedQueue(t *testing.T) {
    65  	fixedSizeMsg := []byte("abcd")
    66  
    67  	conf := Conf{Directory: "/tmp/dq", WriteMmap: true, CustomDecoder: func(ctx context.Context, r *QfileSizeReader) (otherFile bool, data []byte, err error) {
    68  		data = make([]byte, len(fixedSizeMsg))
    69  		err = r.Read(ctx, data)
    70  		return
    71  	}}
    72  	q, err := New(conf)
    73  	assert.Assert(t, err == nil)
    74  	defer func() {
    75  		err = q.Delete()
    76  		assert.Assert(t, err == nil)
    77  	}()
    78  
    79  	// test Read
    80  	n := 1000
    81  	var offsets []int64
    82  	for i := 0; i < n; i++ {
    83  		offset, err := q.Put(fixedSizeMsg)
    84  		assert.Assert(t, err == nil)
    85  		offsets = append(offsets, offset)
    86  
    87  		readData, err := q.Read(nil, offset)
    88  		// fmt.Println(string(readData))
    89  		assert.Assert(t, err == nil && bytes.Equal(readData, fixedSizeMsg), "%v:%v:%v", err, i, offset)
    90  		assert.Assert(t, q.FileMeta(0).MsgCount == uint64(i+1))
    91  	}
    92  
    93  	assert.Assert(t, q.NumFiles() == 1 && q.FileMeta(0).MsgCount == uint64(n))
    94  
    95  	// test StreamOffsetRead
    96  	offsetCh := make(chan int64)
    97  	ch, err := q.StreamOffsetRead(offsetCh)
    98  	go func() {
    99  		for i := 0; i < n; i++ {
   100  			offsetCh <- offsets[i]
   101  		}
   102  	}()
   103  	for i := 0; i < n; i++ {
   104  		streamData, ok := <-ch
   105  		assert.Assert(t, bytes.Equal(streamData.Bytes, fixedSizeMsg) && streamData.Offset == offsets[i], "%v %v", i, ok)
   106  	}
   107  	close(offsetCh)
   108  
   109  	// test StreamRead
   110  	ch, err = q.StreamRead(context.Background(), 0)
   111  	assert.Assert(t, err == nil)
   112  	for i := 0; i < n; i++ {
   113  		streamData := <-ch
   114  		assert.Assert(t, bytes.Equal(streamData.Bytes, fixedSizeMsg))
   115  	}
   116  
   117  	n, err = q.GC()
   118  	assert.Assert(t, err == nil && n == 0)
   119  
   120  }