github.com/nya3jp/tast@v0.0.0-20230601000426-85c8e4d83a9b/src/go.chromium.org/tast/core/rpc/rpc.go (about)

     1  // Copyright 2023 The ChromiumOS Authors
     2  // Use of this source code is governed by a BSD-style license that can be
     3  // found in the LICENSE file.
     4  
     5  // Package rpc provides gRPC utilities for Tast tests.
     6  package rpc
     7  
     8  import (
     9  	"context"
    10  	"os"
    11  	"path/filepath"
    12  
    13  	"google.golang.org/grpc"
    14  
    15  	"go.chromium.org/tast/core/dut"
    16  	"go.chromium.org/tast/core/errors"
    17  	"go.chromium.org/tast/core/internal/protocol"
    18  	"go.chromium.org/tast/core/internal/rpc"
    19  	"go.chromium.org/tast/core/internal/testing"
    20  	"go.chromium.org/tast/core/ssh"
    21  )
    22  
    23  // Client owns a gRPC connection to the DUT for remote tests to use.
    24  type Client struct {
    25  	// Conn is the gRPC connection. Use this to create gRPC service stubs.
    26  	Conn *grpc.ClientConn
    27  
    28  	cl *rpc.SSHClient
    29  }
    30  
    31  // Close closes the connection.
    32  // TODO(b/3042409): Remove ctx param from this method.
    33  func (c *Client) Close(ctx context.Context, opts ...ssh.RunOption) error {
    34  	return c.cl.Close(opts...)
    35  }
    36  
    37  // Dial establishes a gRPC connection to the test bundle executable
    38  // using d and h.
    39  //
    40  // The context passed in must remain valid for as long as the gRPC connection.
    41  // I.e. Don't use the context from within a testing.Poll function.
    42  //
    43  // Example:
    44  //
    45  //	cl, err := rpc.Dial(ctx, d, s.RPCHint())
    46  //	if err != nil {
    47  //		return err
    48  //	}
    49  //	defer cl.Close(ctx)
    50  //
    51  //	fs := base.NewFileSystemClient(cl.Conn)
    52  //
    53  //	res, err := fs.ReadDir(ctx, &base.ReadDirRequest{Dir: "/mnt/stateful_partition"})
    54  //	if err != nil {
    55  //		return err
    56  //	}
    57  func Dial(ctx context.Context, d *dut.DUT, h *testing.RPCHint) (*Client, error) {
    58  	exe, err := os.Executable()
    59  	if err != nil {
    60  		return nil, errors.Wrap(err, "failed to get self bundle name")
    61  	}
    62  
    63  	selfName := filepath.Base(exe)
    64  
    65  	bundlePath := filepath.Join(testing.ExtractLocalBundleDir(h), selfName)
    66  	req := &protocol.HandshakeRequest{
    67  		NeedUserServices: true,
    68  		BundleInitParams: &protocol.BundleInitParams{
    69  			Vars: testing.ExtractTestVars(h),
    70  		},
    71  	}
    72  	cl, err := rpc.DialSSH(ctx, d.Conn(), bundlePath, req, false)
    73  	if err != nil {
    74  		return nil, err
    75  	}
    76  	return &Client{
    77  		Conn: cl.Conn(),
    78  		cl:   cl,
    79  	}, nil
    80  }