wandb

1
package server_test
2

3
import (
4
	"testing"
5

6
	"github.com/wandb/wandb/core/pkg/service"
7
)
8

9
type data struct {
10
	items    map[string]string
11
	step     int64
12
	flush    bool
13
	stepNil  bool
14
	flushNil bool
15
}
16

17
func makeFlushRecord() *service.Record {
18
	record := &service.Record{
19
		RecordType: &service.Record_Request{
20
			Request: &service.Request{
21
				RequestType: &service.Request_Defer{
22
					Defer: &service.DeferRequest{
23
						State: service.DeferRequest_FLUSH_PARTIAL_HISTORY,
24
					},
25
				},
26
			},
27
		},
28
	}
29
	return record
30
}
31

32
func makePartialHistoryRecord(d data) *service.Record {
33
	items := []*service.HistoryItem{}
34
	for k, v := range d.items {
35
		items = append(items, &service.HistoryItem{
36
			Key:       k,
37
			ValueJson: v,
38
		})
39
	}
40
	partialHistoryRequest := &service.PartialHistoryRequest{
41
		Item: items,
42
	}
43
	if !d.stepNil {
44
		partialHistoryRequest.Step = &service.HistoryStep{Num: d.step}
45
	}
46
	if !d.flushNil {
47
		partialHistoryRequest.Action = &service.HistoryAction{Flush: d.flush}
48
	}
49
	record := &service.Record{
50
		RecordType: &service.Record_Request{
51
			Request: &service.Request{
52
				RequestType: &service.Request_PartialHistory{
53
					PartialHistory: partialHistoryRequest,
54
				},
55
			},
56
		},
57
		Control: &service.Control{
58
			MailboxSlot: "junk",
59
		},
60
	}
61
	return record
62
}
63

64
func makeHistoryRecord(d data) *service.Record {
65
	items := []*service.HistoryItem{}
66
	for k, v := range d.items {
67
		items = append(items, &service.HistoryItem{
68
			Key:       k,
69
			ValueJson: v,
70
		})
71
	}
72
	history := &service.HistoryRecord{
73
		Item: items,
74
		Step: &service.HistoryStep{Num: d.step},
75
	}
76
	record := &service.Record{
77
		RecordType: &service.Record_History{
78
			History: history,
79
		},
80
		Control: &service.Control{
81
			MailboxSlot: "junk",
82
		},
83
	}
84
	return record
85
}
86

87
func makeOutput(record *service.Record) data {
88
	switch x := record.GetRecordType().(type) {
89
	case *service.Record_History:
90
		history := x.History
91
		if history == nil {
92
			return data{}
93
		}
94
		items := map[string]string{}
95
		for _, item := range history.Item {
96
			// if strings.HasPrefix(item.Key, "_") {
97
			// 	continue
98
			// }
99
			items[item.Key] = item.ValueJson
100
		}
101
		return data{
102
			items: items,
103
			step:  history.Step.Num,
104
		}
105
	case *service.Record_Request:
106
		state := x.Request.GetDefer().GetState()
107
		if state != service.DeferRequest_FLUSH_PARTIAL_HISTORY {
108
			return data{}
109
		}
110
		return data{
111
			flush: true,
112
		}
113
	default:
114
		return data{}
115
	}
116
}
117

118
type testCase struct {
119
	name     string
120
	input    []data
121
	expected []data
122
}
123

124
func TestHandlePartialHistory(t *testing.T) {
125
	testCases := []testCase{
126
		{
127
			name: "NoFlushIncreaseStepFlush",
128
			input: []data{
129
				{
130
					items: map[string]string{
131
						"key1": "1",
132
						"key2": "2",
133
					},
134
					step:  0,
135
					flush: false,
136
				},
137
				{
138
					items: map[string]string{
139
						"key2": "3",
140
					},
141
					step:  1,
142
					flush: true,
143
				},
144
			},
145
			expected: []data{
146
				{
147
					items: map[string]string{
148
						"key1": "1",
149
						"key2": "2",
150
					},
151
					step: 0,
152
				},
153
				{
154
					items: map[string]string{
155
						"key2": "3",
156
					},
157
					step: 1,
158
				},
159
				{
160
					flush: true,
161
				},
162
			},
163
		},
164
		{
165
			name: "NoFlushNoIncreaseStepFlush",
166
			input: []data{
167
				{
168
					items: map[string]string{
169
						"key1": "1",
170
						"key2": "2",
171
					},
172
					step:  0,
173
					flush: false,
174
				},
175
				{
176
					items: map[string]string{
177
						"key2": "3",
178
					},
179
					step:  0,
180
					flush: true,
181
				},
182
			},
183
			expected: []data{
184
				{
185
					items: map[string]string{
186
						"key1": "1",
187
						"key2": "3",
188
					},
189
					step: 0,
190
				},
191
				{
192
					flush: true,
193
				},
194
			},
195
		},
196
		{
197
			name: "FlushIncreaseStepFlush",
198
			input: []data{
199
				{
200
					items: map[string]string{
201
						"key1": "1",
202
						"key2": "2",
203
					},
204
					step:  0,
205
					flush: true,
206
				},
207
				{
208
					items: map[string]string{
209
						"key2": "3",
210
					},
211
					step:  1,
212
					flush: true,
213
				},
214
			},
215
			expected: []data{
216
				{
217
					items: map[string]string{
218
						"key1": "1",
219
						"key2": "2",
220
					},
221
					step: 0,
222
				},
223
				{
224
					items: map[string]string{
225
						"key2": "3",
226
					},
227
					step: 1,
228
				},
229
				{
230
					flush: true,
231
				},
232
			},
233
		},
234
		{
235
			name: "FlushNoIncreaseStepFlush",
236
			input: []data{
237
				{
238
					items: map[string]string{
239
						"key1": "1",
240
						"key2": "2",
241
					},
242
					step:  0,
243
					flush: true,
244
				},
245
				{
246
					items: map[string]string{
247
						"key2": "3",
248
					},
249
					step:  0,
250
					flush: true,
251
				},
252
			},
253
			expected: []data{
254
				{
255
					items: map[string]string{
256
						"key1": "1",
257
						"key2": "2",
258
					},
259
					step: 0,
260
				},
261
				{
262
					flush: true,
263
				},
264
			},
265
		},
266
		{
267
			name: "FlushIncreaseStepNoFlush",
268
			input: []data{
269
				{
270
					items: map[string]string{
271
						"key1": "1",
272
						"key2": "2",
273
					},
274
					step:  0,
275
					flush: true,
276
				},
277
				{
278
					items: map[string]string{
279
						"key2": "3",
280
					},
281
					step:  1,
282
					flush: false,
283
				},
284
			},
285
			expected: []data{
286
				{
287
					items: map[string]string{
288
						"key1": "1",
289
						"key2": "2",
290
					},
291
					step: 0,
292
				},
293
				{
294
					items: map[string]string{
295
						"key2": "3",
296
					},
297
					step: 1,
298
				},
299
				{
300
					flush: true,
301
				},
302
			},
303
		},
304
		{
305
			name: "FlushNoIncreaseStepNoFlush",
306
			input: []data{
307
				{
308
					items: map[string]string{
309
						"key1": "1",
310
						"key2": "2",
311
					},
312
					step:  0,
313
					flush: true,
314
				},
315
				{
316
					items: map[string]string{
317
						"key2": "3",
318
					},
319
					step:  0,
320
					flush: false,
321
				},
322
			},
323
			expected: []data{
324
				{
325
					items: map[string]string{
326
						"key1": "1",
327
						"key2": "2",
328
					},
329
					step: 0,
330
				},
331
				{
332
					flush: true,
333
				},
334
			},
335
		},
336
		{
337
			name: "NoFlushIncreaseStepNoFlush",
338
			input: []data{
339
				{
340
					items: map[string]string{
341
						"key1": "1",
342
						"key2": "2",
343
					},
344
					step:  0,
345
					flush: false,
346
				},
347
				{
348
					items: map[string]string{
349
						"key2": "3",
350
					},
351
					step:  1,
352
					flush: false,
353
				},
354
			},
355
			expected: []data{
356
				{
357
					items: map[string]string{
358
						"key1": "1",
359
						"key2": "2",
360
					},
361
					step: 0,
362
				},
363
				{
364
					items: map[string]string{
365
						"key2": "3",
366
					},
367
					step: 1,
368
				},
369
				{
370
					flush: true,
371
				},
372
			},
373
		},
374
		{
375
			name: "NoFlushNoIncreaseStepNoFlush",
376
			input: []data{
377
				{
378
					items: map[string]string{
379
						"key1": "1",
380
						"key2": "2",
381
					},
382
					step:  0,
383
					flush: false,
384
				},
385
				{
386
					items: map[string]string{
387
						"key2": "3",
388
					},
389
					step:  0,
390
					flush: false,
391
				},
392
			},
393
			expected: []data{
394
				{
395
					items: map[string]string{
396
						"key1": "1",
397
						"key2": "3",
398
					},
399
					step: 0,
400
				},
401
				{
402
					flush: true,
403
				},
404
			},
405
		},
406
		{
407
			name: "NilStepNilFlushNilStepNilFlush",
408
			input: []data{
409
				{
410
					items: map[string]string{
411
						"key1": "1",
412
						"key2": "2",
413
					},
414
					stepNil:  true,
415
					flushNil: true,
416
				},
417
				{
418
					items: map[string]string{
419
						"key2": "3",
420
					},
421
					stepNil:  true,
422
					flushNil: true,
423
				},
424
			},
425
			expected: []data{
426
				{
427
					items: map[string]string{
428
						"key1": "1",
429
						"key2": "2",
430
					},
431
					step: 0,
432
				},
433
				{
434
					items: map[string]string{
435
						"key2": "3",
436
					},
437
					step: 1,
438
				},
439
				{
440
					flush: true,
441
				},
442
			},
443
		},
444
		{
445
			name: "NilStepNilFlushNilStepNoFlush",
446
			input: []data{
447
				{
448
					items: map[string]string{
449
						"key1": "1",
450
						"key2": "2",
451
					},
452
					stepNil:  true,
453
					flushNil: true,
454
				},
455
				{
456
					items: map[string]string{
457
						"key1": "2",
458
						"key2": "3",
459
					},
460
					stepNil: true,
461
					flush:   false,
462
				},
463
			},
464
			expected: []data{
465
				{
466
					items: map[string]string{
467
						"key1": "1",
468
						"key2": "2",
469
					},
470
					step: 0,
471
				},
472
				{
473
					items: map[string]string{
474
						"key1": "2",
475
						"key2": "3",
476
					},
477
					step: 1,
478
				},
479
				{
480
					flush: true,
481
				},
482
			},
483
		},
484
		{
485
			name: "NilStepNilFlushNilStepFlush",
486
			input: []data{
487
				{
488
					items: map[string]string{
489
						"key1": "1",
490
					},
491
					stepNil:  true,
492
					flushNil: true,
493
				},
494
				{
495
					items: map[string]string{
496
						"key1": "2",
497
					},
498
					stepNil: true,
499
					flush:   true,
500
				},
501
			},
502
			expected: []data{
503
				{
504
					items: map[string]string{
505
						"key1": "1",
506
					},
507
					step: 0,
508
				},
509
				{
510
					items: map[string]string{
511
						"key1": "2",
512
					},
513
					step: 1,
514
				},
515
				{
516
					flush: true,
517
				},
518
			},
519
		},
520
		{
521
			name: "StepNoFlushNilStepNilFlush",
522
			input: []data{
523
				{
524
					items: map[string]string{
525
						"key1": "1",
526
					},
527
					step:  1,
528
					flush: false,
529
				},
530
				{
531
					items: map[string]string{
532
						"key1": "2",
533
					},
534
					stepNil:  true,
535
					flushNil: true,
536
				},
537
			},
538
			expected: []data{
539
				{
540
					items: map[string]string{
541
						"key1": "2",
542
					},
543
					step: 1,
544
				},
545
				{
546
					flush: true,
547
				},
548
			},
549
		},
550
		{
551
			name: "StepFlushNilStepNilFlush",
552
			input: []data{
553
				{
554
					items: map[string]string{
555
						"key1": "1",
556
					},
557
					step:  1,
558
					flush: true,
559
				},
560
				{
561
					items: map[string]string{
562
						"key1": "2",
563
					},
564
					stepNil:  true,
565
					flushNil: true,
566
				},
567
			},
568
			expected: []data{
569
				{
570
					items: map[string]string{
571
						"key1": "1",
572
					},
573
					step: 1,
574
				},
575
				{
576
					items: map[string]string{
577
						"key1": "2",
578
					},
579
					step: 2,
580
				},
581
				{
582
					flush: true,
583
				},
584
			},
585
		},
586
		{
587
			name: "StepNilFlushNilStepNilFlush",
588
			input: []data{
589
				{
590
					items: map[string]string{
591
						"key1": "1",
592
					},
593
					step:     1,
594
					flushNil: true,
595
				},
596
				{
597
					items: map[string]string{
598
						"key1": "2",
599
					},
600
					stepNil:  true,
601
					flushNil: true,
602
				},
603
			},
604
			expected: []data{
605
				{
606
					items: map[string]string{
607
						"key1": "2",
608
					},
609
					step: 1,
610
				},
611
				{
612
					flush: true,
613
				},
614
			},
615
		},
616
		{
617
			name: "NilStepFlushNilStepNilFlush",
618
			input: []data{
619
				{
620
					items: map[string]string{
621
						"key1": "1",
622
					},
623
					stepNil: true,
624
					flush:   true,
625
				},
626
				{
627
					items: map[string]string{
628
						"key1": "2",
629
					},
630
					stepNil:  true,
631
					flushNil: true,
632
				},
633
			},
634
			expected: []data{
635
				{
636
					items: map[string]string{
637
						"key1": "1",
638
					},
639
					step: 0,
640
				},
641
				{
642
					items: map[string]string{
643
						"key1": "2",
644
					},
645
					step: 1,
646
				},
647
				{
648
					flush: true,
649
				},
650
			},
651
		},
652
		{
653
			name: "NilStepNoFlushNilStepNilFlush",
654
			input: []data{
655
				{
656
					items: map[string]string{
657
						"key1": "1",
658
					},
659
					stepNil: true,
660
					flush:   false,
661
				},
662
				{
663
					items: map[string]string{
664
						"key1": "2",
665
					},
666
					stepNil:  true,
667
					flushNil: true,
668
				},
669
			},
670
			expected: []data{
671
				{
672
					items: map[string]string{
673
						"key1": "2",
674
					},
675
					step: 0,
676
				},
677
				{
678
					flush: true,
679
				},
680
			},
681
		},
682
	}
683

684
	for _, tc := range testCases {
685
		t.Run(tc.name, func(t *testing.T) {
686
			inChan, loopbackChan := makeInboundChannels()
687
			fwdChan, outChan := makeOutboundChannels()
688

689
			makeHandler(inChan, loopbackChan, fwdChan, outChan, false)
690

691
			for _, d := range tc.input {
692
				record := makePartialHistoryRecord(d)
693
				inChan <- record
694
			}
695

696
			inChan <- makeFlushRecord()
697

698
			for _, d := range tc.expected {
699
				record := <-fwdChan
700
				actual := makeOutput(record)
701
				if actual.step != d.step {
702
					t.Errorf("expected step %v, got %v", d.step, actual.step)
703
				}
704
				for k, v := range d.items {
705
					if actual.items[k] != v {
706
						t.Errorf("expected %v, got %v", v, actual.items[k])
707
					}
708
				}
709
				if d.flush != actual.flush {
710
					t.Errorf("expected %v, got %v", d.flush, d.flush)
711
				}
712
			}
713
		},
714
		)
715
	}
716
}
717

718
func TestHandleHistory(t *testing.T) {
719
	testCases := []testCase{
720
		{
721
			name: "IncreaseStep",
722
			input: []data{
723
				{
724
					items: map[string]string{
725
						"key1": "1",
726
						"key2": "2",
727
					},
728
					step: 0,
729
				},
730
				{
731
					items: map[string]string{
732
						"key2": "3",
733
					},
734
					step: 1,
735
				},
736
			},
737
			expected: []data{
738
				{
739
					items: map[string]string{
740
						"key1":     "1",
741
						"key2":     "2",
742
						"_step":    "0",
743
						"_runtime": "0.000000",
744
					},
745
					step: 0,
746
				},
747
				{
748
					items: map[string]string{
749
						"key2":     "3",
750
						"_step":    "1",
751
						"_runtime": "0.000000",
752
					},
753
					step: 1,
754
				},
755
				{
756
					flush: true,
757
				},
758
			},
759
		},
760
		// { // TODO: mock out time
761
		// 	name: "Timestamp",
762
		// 	input: []data{
763
		// 		{
764
		// 			items: map[string]string{
765
		// 				"key1":       "1",
766
		// 				"key2":       "2",
767
		// 				"_timestamp": "1.257894e+09",
768
		// 			},
769
		// 			step: 0,
770
		// 		},
771
		// 	},
772
		// 	expected: []data{
773
		// 		{
774
		// 			items: map[string]string{
775
		// 				"key1":     "1",
776
		// 				"key2":     "2",
777
		// 				"_runtime": "63393490800.000000",
778
		// 				"_step":    "0",
779
		// 			},
780
		// 			step: 0,
781
		// 		},
782
		// 		{
783
		// 			flush: true,
784
		// 		},
785
		// 	},
786
		// },
787
	}
788
	for _, tc := range testCases {
789
		t.Run(tc.name, func(t *testing.T) {
790
			inChan, loopbackChan := makeInboundChannels()
791
			fwdChan, outChan := makeOutboundChannels()
792

793
			makeHandler(inChan, loopbackChan, fwdChan, outChan, false)
794

795
			for _, d := range tc.input {
796
				record := makeHistoryRecord(d)
797
				inChan <- record
798
			}
799

800
			inChan <- makeFlushRecord()
801

802
			for _, d := range tc.expected {
803
				record := <-fwdChan
804
				actual := makeOutput(record)
805
				if actual.step != d.step {
806
					t.Errorf("expected step %v, got %v", d.step, actual.step)
807
				}
808
				for k, v := range actual.items {
809
					if d.items[k] != v {
810
						t.Errorf("expected %v, got %v", v, actual.items[k])
811
					}
812
				}
813
				if d.flush != actual.flush {
814
					t.Errorf("expected %v, got %v", d.flush, d.flush)
815
				}
816
			}
817
		})
818
	}
819

820
}
821