sdk/go/arvadostest/test_cluster.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package arvadostest
   6
   7 import (
   8         "context"
   9         "net/url"
  10
  11         "git.arvados.org/arvados.git/lib/boot"
  12         "git.arvados.org/arvados.git/lib/controller/rpc"
  13         "git.arvados.org/arvados.git/lib/service"
  14         "git.arvados.org/arvados.git/sdk/go/arvados"
  15         "git.arvados.org/arvados.git/sdk/go/arvadosclient"
  16         "git.arvados.org/arvados.git/sdk/go/auth"
  17         "git.arvados.org/arvados.git/sdk/go/ctxlog"
  18         "git.arvados.org/arvados.git/sdk/go/keepclient"
  19         "gopkg.in/check.v1"
  20 )
  21
  22 // TestCluster stores a working test cluster data
  23 type TestCluster struct {
  24         Super         boot.Supervisor
  25         Config        arvados.Config
  26         ControllerURL *url.URL
  27         ClusterID     string
  28 }
  29
  30 type logger struct {
  31         loggerfunc func(...interface{})
  32 }
  33
  34 func (l logger) Log(args ...interface{}) {
  35         l.loggerfunc(args)
  36 }
  37
  38 // NewTestCluster loads the provided configuration, and sets up a test cluster
  39 // ready for being started.
  40 func NewTestCluster(srcPath, clusterID string, cfg *arvados.Config, listenHost string, logWriter func(...interface{})) *TestCluster {
  41         return &TestCluster{
  42                 Super: boot.Supervisor{
  43                         SourcePath:           srcPath,
  44                         ClusterType:          "test",
  45                         ListenHost:           listenHost,
  46                         ControllerAddr:       ":0",
  47                         OwnTemporaryDatabase: true,
  48                         Stderr: &service.LogPrefixer{
  49                                 Writer: ctxlog.LogWriter(logWriter),
  50                                 Prefix: []byte("[" + clusterID + "] ")},
  51                 },
  52                 Config:    *cfg,
  53                 ClusterID: clusterID,
  54         }
  55 }
  56
  57 // Start the test cluster.
  58 func (tc *TestCluster) Start() {
  59         tc.Super.Start(context.Background(), &tc.Config, "-")
  60 }
  61
  62 // WaitReady waits for all components to report healthy, and finishes setting
  63 // up the TestCluster struct.
  64 func (tc *TestCluster) WaitReady() bool {
  65         au, ok := tc.Super.WaitReady()
  66         if !ok {
  67                 return ok
  68         }
  69         u := url.URL(*au)
  70         tc.ControllerURL = &u
  71         return ok
  72 }
  73
  74 // ClientsWithToken returns Context, Arvados.Client and keepclient structs
  75 // initialized to connect to the cluster with the supplied Arvados token.
  76 func (tc *TestCluster) ClientsWithToken(token string) (context.Context, *arvados.Client, *keepclient.KeepClient) {
  77         cl := tc.Config.Clusters[tc.ClusterID]
  78         ctx := auth.NewContext(context.Background(), auth.NewCredentials(token))
  79         ac, err := arvados.NewClientFromConfig(&cl)
  80         if err != nil {
  81                 panic(err)
  82         }
  83         ac.AuthToken = token
  84         arv, err := arvadosclient.New(ac)
  85         if err != nil {
  86                 panic(err)
  87         }
  88         kc := keepclient.New(arv)
  89         return ctx, ac, kc
  90 }
  91
  92 // UserClients logs in as a user called "example", get the user's API token,
  93 // initialize clients with the API token, set up the user and
  94 // optionally activate the user.  Return client structs for
  95 // communicating with the cluster on behalf of the 'example' user.
  96 func (tc *TestCluster) UserClients(rootctx context.Context, c *check.C, conn *rpc.Conn, authEmail string, activate bool) (context.Context, *arvados.Client, *keepclient.KeepClient, arvados.User) {
  97         login, err := conn.UserSessionCreate(rootctx, rpc.UserSessionCreateOptions{
  98                 ReturnTo: ",https://example.com",
  99                 AuthInfo: rpc.UserSessionAuthInfo{
 100                         Email:     authEmail,
 101                         FirstName: "Example",
 102                         LastName:  "User",
 103                         Username:  "example",
 104                 },
 105         })
 106         c.Assert(err, check.IsNil)
 107         redirURL, err := url.Parse(login.RedirectLocation)
 108         c.Assert(err, check.IsNil)
 109         userToken := redirURL.Query().Get("api_token")
 110         c.Logf("user token: %q", userToken)
 111         ctx, ac, kc := tc.ClientsWithToken(userToken)
 112         user, err := conn.UserGetCurrent(ctx, arvados.GetOptions{})
 113         c.Assert(err, check.IsNil)
 114         _, err = conn.UserSetup(rootctx, arvados.UserSetupOptions{UUID: user.UUID})
 115         c.Assert(err, check.IsNil)
 116         if activate {
 117                 _, err = conn.UserActivate(rootctx, arvados.UserActivateOptions{UUID: user.UUID})
 118                 c.Assert(err, check.IsNil)
 119                 user, err = conn.UserGetCurrent(ctx, arvados.GetOptions{})
 120                 c.Assert(err, check.IsNil)
 121                 c.Logf("user UUID: %q", user.UUID)
 122                 if !user.IsActive {
 123                         c.Fatalf("failed to activate user -- %#v", user)
 124                 }
 125         }
 126         return ctx, ac, kc, user
 127 }
 128
 129 // RootClients returns Context, arvados.Client and keepclient structs initialized
 130 // to communicate with the cluster as the system root user.
 131 func (tc *TestCluster) RootClients() (context.Context, *arvados.Client, *keepclient.KeepClient) {
 132         return tc.ClientsWithToken(tc.Config.Clusters[tc.ClusterID].SystemRootToken)
 133 }
 134
 135 // AnonymousClients returns Context, arvados.Client and keepclient structs initialized
 136 // to communicate with the cluster as the anonymous user.
 137 func (tc *TestCluster) AnonymousClients() (context.Context, *arvados.Client, *keepclient.KeepClient) {
 138         return tc.ClientsWithToken(tc.Config.Clusters[tc.ClusterID].Users.AnonymousUserToken)
 139 }
 140
 141 // Conn gets rpc connection struct initialized to communicate with the
 142 // specified cluster.
 143 func (tc *TestCluster) Conn() *rpc.Conn {
 144         return rpc.NewConn(tc.ClusterID, tc.ControllerURL, true, rpc.PassthroughTokenProvider)
 145 }