wandb

1
package server_test
2

3
import (
4
	"context"
5
	"testing"
6

7
	"github.com/Khan/genqlient/graphql"
8
	"github.com/golang/mock/gomock"
9
	"github.com/stretchr/testify/assert"
10
	"github.com/wandb/wandb/core/internal/coretest"
11
	"github.com/wandb/wandb/core/internal/gql"
12
	"github.com/wandb/wandb/core/pkg/observability"
13
	"github.com/wandb/wandb/core/pkg/server"
14
	"github.com/wandb/wandb/core/pkg/service"
15
	"google.golang.org/protobuf/types/known/wrapperspb"
16
)
17

18
func makeSender(client graphql.Client, resultChan chan *service.Result) *server.Sender {
19
	ctx, cancel := context.WithCancel(context.Background())
20
	logger := observability.NewNoOpLogger()
21
	sender := server.NewSender(
22
		ctx,
23
		cancel,
24
		logger,
25
		&service.Settings{
26
			RunId: &wrapperspb.StringValue{Value: "run1"},
27
		},
28
		server.WithSenderFwdChannel(make(chan *service.Record, 1)),
29
		server.WithSenderOutChannel(resultChan),
30
	)
31
	sender.SetGraphqlClient(client)
32
	return sender
33
}
34

35
func TestSendRun(t *testing.T) {
36
	// Verify that project and entity are properly passed through to graphql
37
	to := coretest.MakeTestObject(t)
38
	defer to.TeardownTest()
39

40
	sender := makeSender(to.MockClient, make(chan *service.Result, 1))
41

42
	run := &service.Record{
43
		RecordType: &service.Record_Run{
44
			Run: &service.RunRecord{
45
				Config:  to.MakeConfig(),
46
				Project: "testProject",
47
				Entity:  "testEntity",
48
			}},
49
		Control: &service.Control{
50
			MailboxSlot: "junk",
51
		},
52
	}
53

54
	respEncode := &graphql.Response{
55
		Data: &gql.UpsertBucketResponse{
56
			UpsertBucket: &gql.UpsertBucketUpsertBucketUpsertBucketPayload{
57
				Bucket: &gql.UpsertBucketUpsertBucketUpsertBucketPayloadBucketRun{
58
					DisplayName: coretest.StrPtr("FakeName"),
59
					Project: &gql.UpsertBucketUpsertBucketUpsertBucketPayloadBucketRunProject{
60
						Name: "FakeProject",
61
						Entity: gql.UpsertBucketUpsertBucketUpsertBucketPayloadBucketRunProjectEntity{
62
							Name: "FakeEntity",
63
						},
64
					},
65
				},
66
			},
67
		},
68
	}
69

70
	to.MockClient.EXPECT().MakeRequest(
71
		gomock.Any(), // context.Context
72
		gomock.Any(), // *graphql.Request
73
		gomock.Any(), // *graphql.Response
74
	).Return(nil).Do(coretest.InjectResponse(
75
		respEncode,
76
		func(vars coretest.RequestVars) {
77
			assert.Equal(t, "testEntity", vars["entity"])
78
			assert.Equal(t, "testProject", vars["project"])
79
		},
80
	))
81

82
	sender.SendRecord(run)
83
	<-sender.GetOutboundChannel()
84
}
85

86
func TestSendLinkArtifact(t *testing.T) {
87
	// Verify that arguments are properly passed through to graphql
88
	to := coretest.MakeTestObject(t)
89
	defer to.TeardownTest()
90

91
	sender := makeSender(to.MockClient, make(chan *service.Result, 1))
92

93
	respEncode := &graphql.Response{
94
		Data: &gql.LinkArtifactResponse{
95
			LinkArtifact: &gql.LinkArtifactLinkArtifactLinkArtifactPayload{
96
				VersionIndex: coretest.IntPtr(0),
97
			},
98
		}}
99

100
	// 1. When both clientId and serverId are sent, serverId is used
101
	linkArtifact := &service.Record{
102
		RecordType: &service.Record_LinkArtifact{
103
			LinkArtifact: &service.LinkArtifactRecord{
104
				ClientId:         "clientId",
105
				ServerId:         "serverId",
106
				PortfolioName:    "portfolioName",
107
				PortfolioEntity:  "portfolioEntity",
108
				PortfolioProject: "portfolioProject",
109
			}},
110
		Control: &service.Control{
111
			MailboxSlot: "junk",
112
		},
113
	}
114

115
	to.MockClient.EXPECT().MakeRequest(
116
		gomock.Any(), // context.Context
117
		gomock.Any(), // *graphql.Request
118
		gomock.Any(), // *graphql.Response
119
	).Return(nil).Do(coretest.InjectResponse(
120
		respEncode,
121
		func(vars coretest.RequestVars) {
122
			assert.Equal(t, "portfolioProject", vars["projectName"])
123
			assert.Equal(t, "portfolioEntity", vars["entityName"])
124
			assert.Equal(t, "portfolioName", vars["artifactPortfolioName"])
125
			assert.Nil(t, vars["clientId"])
126
			assert.Equal(t, "serverId", vars["artifactId"])
127
		},
128
	))
129

130
	sender.SendRecord(linkArtifact)
131
	<-sender.GetOutboundChannel()
132

133
	// 2. When only clientId is sent, clientId is used
134
	linkArtifact = &service.Record{
135
		RecordType: &service.Record_LinkArtifact{
136
			LinkArtifact: &service.LinkArtifactRecord{
137
				ClientId:         "clientId",
138
				ServerId:         "",
139
				PortfolioName:    "portfolioName",
140
				PortfolioEntity:  "portfolioEntity",
141
				PortfolioProject: "portfolioProject",
142
			}},
143
		Control: &service.Control{
144
			MailboxSlot: "junk",
145
		},
146
	}
147

148
	to.MockClient.EXPECT().MakeRequest(
149
		gomock.Any(), // context.Context
150
		gomock.Any(), // *graphql.Request
151
		gomock.Any(), // *graphql.Response
152
	).Return(nil).Do(coretest.InjectResponse(
153
		respEncode,
154
		func(vars coretest.RequestVars) {
155
			assert.Equal(t, "portfolioProject", vars["projectName"])
156
			assert.Equal(t, "portfolioEntity", vars["entityName"])
157
			assert.Equal(t, "portfolioName", vars["artifactPortfolioName"])
158
			assert.Equal(t, "clientId", vars["clientId"])
159
			assert.Nil(t, vars["artifactId"])
160
		},
161
	))
162

163
	sender.SendRecord(linkArtifact)
164
	<-sender.GetOutboundChannel()
165

166
	// 2. When only serverId is sent, serverId is used
167
	linkArtifact = &service.Record{
168
		RecordType: &service.Record_LinkArtifact{
169
			LinkArtifact: &service.LinkArtifactRecord{
170
				ClientId:         "",
171
				ServerId:         "serverId",
172
				PortfolioName:    "portfolioName",
173
				PortfolioEntity:  "portfolioEntity",
174
				PortfolioProject: "portfolioProject",
175
			}},
176
		Control: &service.Control{
177
			MailboxSlot: "junk",
178
		},
179
	}
180

181
	to.MockClient.EXPECT().MakeRequest(
182
		gomock.Any(), // context.Context
183
		gomock.Any(), // *graphql.Request
184
		gomock.Any(), // *graphql.Response
185
	).Return(nil).Do(coretest.InjectResponse(
186
		respEncode,
187
		func(vars coretest.RequestVars) {
188
			assert.Equal(t, "portfolioProject", vars["projectName"])
189
			assert.Equal(t, "portfolioEntity", vars["entityName"])
190
			assert.Equal(t, "portfolioName", vars["artifactPortfolioName"])
191
			assert.Nil(t, vars["clientId"])
192
			assert.Equal(t, "serverId", vars["artifactId"])
193
		},
194
	))
195

196
	sender.SendRecord(linkArtifact)
197
	<-sender.GetOutboundChannel()
198
}
199

200
func TestSendUseArtifact(t *testing.T) {
201
	to := coretest.MakeTestObject(t)
202
	defer to.TeardownTest()
203

204
	sender := makeSender(to.MockClient, make(chan *service.Result, 1))
205

206
	useArtifact := &service.Record{
207
		RecordType: &service.Record_UseArtifact{
208
			UseArtifact: &service.UseArtifactRecord{
209
				Id:      "artifactId",
210
				Type:    "job",
211
				Name:    "artifactName",
212
				Partial: nil,
213
			},
214
		},
215
	}
216
	// verify doesn't panic if used job artifact
217
	sender.SendRecord(useArtifact)
218

219
	// verify doesn't panic if partial job is broken
220
	useArtifact = &service.Record{
221
		RecordType: &service.Record_UseArtifact{
222
			UseArtifact: &service.UseArtifactRecord{
223
				Id:   "artifactId",
224
				Type: "job",
225
				Name: "artifactName",
226
				Partial: &service.PartialJobArtifact{
227
					JobName: "jobName",
228
					SourceInfo: &service.JobSource{
229
						SourceType: "repo",
230
						Source: &service.Source{
231
							Git: &service.GitSource{
232
								GitInfo: &service.GitInfo{
233
									Commit: "commit",
234
									Remote: "remote",
235
								},
236
							},
237
						},
238
					},
239
				},
240
			},
241
		},
242
	}
243
	sender.SendRecord(useArtifact)
244
}
245

246
func TestSendArtifact(t *testing.T) {
247
	// Verify that arguments are properly passed through to graphql
248
	to := coretest.MakeTestObject(t)
249
	defer to.TeardownTest()
250

251
	sender := makeSender(to.MockClient, make(chan *service.Result, 1))
252

253
	// 1. When both clientId and serverId are sent, serverId is used
254
	artifact := &service.Record{
255
		RecordType: &service.Record_Artifact{
256
			Artifact: &service.ArtifactRecord{
257
				RunId:   "test-run-id",
258
				Project: "test-project",
259
				Entity:  "test-entity",
260
				Type:    "test-type",
261
				Name:    "test-artifact",
262
				Digest:  "test-digest",
263
				Aliases: []string{"latest"},
264
				Manifest: &service.ArtifactManifest{
265
					Version:       1,
266
					StoragePolicy: "wandb-storage-policy-v1",
267
					Contents: []*service.ArtifactManifestEntry{{
268
						Path:      "test1",
269
						Digest:    "test1-digest",
270
						Size:      1,
271
						LocalPath: "/test/local/path",
272
					},
273
					},
274
				},
275
				Finalize:         true,
276
				ClientId:         "client-id",
277
				SequenceClientId: "sequence-client-id",
278
			}},
279
	}
280
	createArtifactRespEncode := &graphql.Response{
281
		Data: &gql.CreateArtifactResponse{
282
			CreateArtifact: &gql.CreateArtifactCreateArtifactCreateArtifactPayload{
283
				Artifact: gql.CreateArtifactCreateArtifactCreateArtifactPayloadArtifact{
284
					Id: "artifact-id",
285
				},
286
			},
287
		}}
288
	to.MockClient.EXPECT().MakeRequest(
289
		gomock.Any(), // context.Context
290
		gomock.Any(), // *graphql.Request
291
		gomock.Any(), // *graphql.Response
292
	).Return(nil).Do(coretest.InjectResponse(
293
		createArtifactRespEncode,
294
		func(vars coretest.RequestVars) {
295
			assert.Equal(t, "test-entity", vars["entityName"])
296
		},
297
	))
298
	sender.SendRecord(artifact)
299
}
300