github.com/ethersphere/bee/v2@v2.2.0/pkg/storer/subscribe_push_test.go (about)

     1  // Copyright 2023 The Swarm Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package storer_test
     6  
     7  import (
     8  	"bytes"
     9  	"context"
    10  	"sync"
    11  	"testing"
    12  	"time"
    13  
    14  	chunktesting "github.com/ethersphere/bee/v2/pkg/storage/testing"
    15  	storer "github.com/ethersphere/bee/v2/pkg/storer"
    16  	"github.com/ethersphere/bee/v2/pkg/swarm"
    17  )
    18  
    19  func TestPushSubscriber(t *testing.T) {
    20  	t.Parallel()
    21  
    22  	baseAddr := swarm.RandAddress(t)
    23  	t.Run("inmem", func(t *testing.T) {
    24  		t.Parallel()
    25  		memStorer := memStorer(t, dbTestOps(baseAddr, 10, nil, nil, time.Second))
    26  		testPushSubscriber(t, memStorer)
    27  	})
    28  	t.Run("disk", func(t *testing.T) {
    29  		t.Parallel()
    30  		diskStorer := diskStorer(t, dbTestOps(baseAddr, 0, nil, nil, time.Second))
    31  		testPushSubscriber(t, diskStorer)
    32  	})
    33  }
    34  
    35  func testPushSubscriber(t *testing.T, newLocalstore func() (*storer.DB, error)) {
    36  	t.Helper()
    37  
    38  	lstore, err := newLocalstore()
    39  	if err != nil {
    40  		t.Fatal(err)
    41  	}
    42  
    43  	chunks := chunktesting.GenerateTestRandomChunks(18)
    44  	chunksProcessedMu := sync.Mutex{}
    45  	chunksProcessed := make(map[string]struct{})
    46  
    47  	uploadChunks := func(chunksToUpload ...swarm.Chunk) {
    48  		id, err := lstore.NewSession()
    49  		if err != nil {
    50  			t.Fatal(err)
    51  		}
    52  
    53  		p, err := lstore.Upload(context.Background(), false, id.TagID)
    54  		if err != nil {
    55  			t.Fatal(err)
    56  		}
    57  
    58  		for _, ch := range chunksToUpload {
    59  			if err := p.Put(context.Background(), ch); err != nil {
    60  				t.Fatal(err)
    61  			}
    62  		}
    63  
    64  		err = p.Done(chunksToUpload[len(chunksToUpload)-1].Address())
    65  		if err != nil {
    66  			t.Fatalf("session.Done(...): unexpected error: %v", err)
    67  		}
    68  	}
    69  
    70  	verify := func(got swarm.Chunk) {
    71  		idx := swarm.IndexOfChunkWithAddress(chunks, got.Address())
    72  		if idx < 0 {
    73  			t.Fatalf("chunk not found %s", got.Address())
    74  		}
    75  		want := chunks[idx]
    76  		if !got.Equal(want) {
    77  			t.Fatalf("chunk not equal %s", got.Address())
    78  		}
    79  		gotStamp, err := got.Stamp().MarshalBinary()
    80  		if err != nil {
    81  			t.Fatalf("stamp.MarshalBinary: unexpected error: %v", err)
    82  		}
    83  		wantStamp, err := want.Stamp().MarshalBinary()
    84  		if err != nil {
    85  			t.Fatalf("stamp.MarshalBinary: unexpected error: %v", err)
    86  		}
    87  		if !bytes.Equal(gotStamp, wantStamp) {
    88  			t.Fatalf("stamps dont match")
    89  		}
    90  
    91  		chunksProcessedMu.Lock()
    92  		chunksProcessed[got.Address().ByteString()] = struct{}{}
    93  		chunksProcessedMu.Unlock()
    94  	}
    95  
    96  	// prepopulate database with some chunks
    97  	// before the subscription
    98  	uploadChunks(chunks[:10]...)
    99  
   100  	// set a timeout on subscription
   101  	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
   102  	defer cancel()
   103  
   104  	ch, stop := lstore.SubscribePush(ctx)
   105  	defer stop()
   106  
   107  	// receive and validate addresses from the subscription
   108  	done := make(chan struct{})
   109  	go func() {
   110  		defer close(done)
   111  		for {
   112  			select {
   113  			case got, ok := <-ch:
   114  				if !ok {
   115  					return
   116  				}
   117  				verify(got)
   118  
   119  				chunksProcessedMu.Lock()
   120  				count := len(chunksProcessed)
   121  				chunksProcessedMu.Unlock()
   122  
   123  				if count == len(chunks) {
   124  					return
   125  				}
   126  			case <-ctx.Done():
   127  				return
   128  			}
   129  		}
   130  	}()
   131  
   132  	// upload some chunks just after subscribe
   133  	uploadChunks(chunks[10:15]...)
   134  
   135  	time.Sleep(200 * time.Millisecond)
   136  
   137  	// upload some chunks after some short time
   138  	// to ensure that subscription will include them
   139  	// in a dynamic environment
   140  	uploadChunks(chunks[15:]...)
   141  
   142  	<-done
   143  }