go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/resultdb/internal/artifactcontent/testutil/test_util.go (about)

     1  // Copyright 2021 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package testutil
    16  
    17  import (
    18  	"context"
    19  	"io"
    20  
    21  	"google.golang.org/genproto/googleapis/bytestream"
    22  	"google.golang.org/grpc"
    23  )
    24  
    25  type FakeCASReader struct {
    26  	grpc.ClientStream // implements the rest of grpc.ClientStream
    27  
    28  	Res         []*bytestream.ReadResponse
    29  	ResIndex    int
    30  	ResDataPos  int
    31  	ResErr      error
    32  	ResErrIndex int
    33  
    34  	ReadLimit int
    35  	nRead     int
    36  }
    37  
    38  func (r *FakeCASReader) Recv() (*bytestream.ReadResponse, error) {
    39  	if r.ResErr != nil && r.ResErrIndex == r.ResIndex {
    40  		return nil, r.ResErr
    41  	}
    42  
    43  	limitAvail := r.ReadLimit - r.nRead
    44  	if r.ReadLimit > 0 && limitAvail == 0 {
    45  		return nil, io.EOF
    46  	}
    47  
    48  	if r.ResIndex < len(r.Res) {
    49  		// calculate how much data can be read further from the current Res.
    50  		res := r.Res[r.ResIndex]
    51  		nRead := len(res.Data) - r.ResDataPos
    52  		if r.ReadLimit != 0 && nRead > limitAvail {
    53  			nRead = limitAvail
    54  		}
    55  
    56  		data := res.Data[r.ResDataPos : r.ResDataPos+nRead]
    57  		r.nRead += nRead
    58  		r.ResDataPos += nRead
    59  		if r.ResDataPos == len(res.Data) {
    60  			r.ResDataPos = 0
    61  			r.ResIndex++
    62  		}
    63  		return &bytestream.ReadResponse{Data: data}, nil
    64  	}
    65  	return nil, io.EOF
    66  }
    67  
    68  type FakeByteStreamClient struct {
    69  	ExtraResponseData []byte
    70  }
    71  
    72  func (c *FakeByteStreamClient) Read(ctx context.Context, in *bytestream.ReadRequest, opts ...grpc.CallOption) (bytestream.ByteStream_ReadClient, error) {
    73  	res := []*bytestream.ReadResponse{
    74  		{Data: []byte("contentspart1\n")},
    75  	}
    76  	if len(c.ExtraResponseData) > 0 {
    77  		res = append(res, &bytestream.ReadResponse{Data: c.ExtraResponseData})
    78  	}
    79  	return &FakeCASReader{
    80  		Res: res,
    81  	}, nil
    82  }
    83  
    84  func (c *FakeByteStreamClient) Write(ctx context.Context, opts ...grpc.CallOption) (bytestream.ByteStream_WriteClient, error) {
    85  	return nil, nil
    86  }
    87  
    88  func (c *FakeByteStreamClient) QueryWriteStatus(ctx context.Context, in *bytestream.QueryWriteStatusRequest, opts ...grpc.CallOption) (*bytestream.QueryWriteStatusResponse, error) {
    89  	return nil, nil
    90  }