github.com/bazelbuild/remote-apis-sdks@v0.0.0-20240425170053-8a36686a6350/go/pkg/fakes/logstreams.go (about)

     1  package fakes
     2  
     3  import (
     4  	"fmt"
     5  	"strings"
     6  
     7  	"google.golang.org/grpc/codes"
     8  	"google.golang.org/grpc/status"
     9  
    10  	bsgrpc "google.golang.org/genproto/googleapis/bytestream"
    11  	bspb "google.golang.org/genproto/googleapis/bytestream"
    12  )
    13  
    14  // LogStreams is a fake logstream implementation that implements the bytestream Read command.
    15  type LogStreams struct {
    16  	// streams is a map containing the logstreams. Each logstream consists of a list of chunks. When
    17  	// read, the Read() method will send each chunk one a time.
    18  	streams map[string][][]byte
    19  }
    20  
    21  // NewLogStreams returns a new empty fake logstream implementation.
    22  func NewLogStreams() *LogStreams {
    23  	return &LogStreams{streams: make(map[string][][]byte)}
    24  }
    25  
    26  // Clear removes all logstreams.
    27  func (l *LogStreams) Clear() {
    28  	l.streams = make(map[string][][]byte)
    29  }
    30  
    31  // Put stores a new logstream.
    32  func (l *LogStreams) Put(name string, chunks ...string) error {
    33  	if _, ok := l.streams[name]; ok {
    34  		return fmt.Errorf("stream with name %q already exists", name)
    35  	}
    36  	var bs [][]byte
    37  	for _, chunk := range chunks {
    38  		bs = append(bs, []byte(chunk))
    39  	}
    40  	l.streams[name] = bs
    41  	return nil
    42  }
    43  
    44  // Read implements the Bytestream Read command. The chunks of the requested logstream are sent one
    45  // at a time.
    46  func (l *LogStreams) Read(req *bspb.ReadRequest, stream bsgrpc.ByteStream_ReadServer) error {
    47  	path := strings.Split(req.ResourceName, "/")
    48  	if len(path) != 3 || path[0] != "instance" || path[1] != "logstreams" {
    49  		return status.Error(codes.InvalidArgument, "test fake expected resource name of the form \"instance/logstreams/<name>\"")
    50  	}
    51  
    52  	chunks, ok := l.streams[path[2]]
    53  	if !ok {
    54  		return status.Error(codes.NotFound, "logstream not found")
    55  	}
    56  
    57  	for _, chunk := range chunks {
    58  		if err := stream.Send(&bspb.ReadResponse{Data: chunk}); err != nil {
    59  			return err
    60  		}
    61  	}
    62  	return nil
    63  }