66package cli
77
88import (
9+ "archive/zip"
910 "compress/gzip"
1011 "encoding/csv"
1112 "encoding/gob"
@@ -74,8 +75,11 @@ func TestTSDumpUploadE2E(t *testing.T) {
7475 debugTimeSeriesDumpOpts .zendeskTicket = "zd-test"
7576 debugTimeSeriesDumpOpts .organizationName = "test-org"
7677 debugTimeSeriesDumpOpts .userName = "test-user"
77- dumpFilePath := generateMockTSDumpFromCSV (t , d .Input )
7878
79+ var compression string
80+ d .MaybeScanArgs (t , "compression" , & compression )
81+
82+ dumpFilePath := generateMockTSDumpFromCSV (t , d .Input , withCompression (compression ))
7983 var clusterLabel , apiKey string
8084 if d .HasArg ("cluster-label" ) {
8185 d .ScanArgs (t , "cluster-label" , & clusterLabel )
@@ -276,7 +280,7 @@ func TestTSDumpUploadWithEmbeddedMetadataDataDriven(t *testing.T) {
276280 case "upload-datadog-embedded-only" :
277281 // Embedded metadata only (no YAML file)
278282 // This tests that store metrics get proper node_id tags based on embedded store-to-node mapping
279- dumpFilePath = generateMockTSDumpWithEmbeddedMetadata (t , d .Input )
283+ dumpFilePath = generateMockTSDumpFromCSV (t , d .Input , withEmbeddedMetadata () )
280284
281285 case "upload-datadog-yaml-only" :
282286 // YAML file only (no embedded metadata)
@@ -292,7 +296,7 @@ func TestTSDumpUploadWithEmbeddedMetadataDataDriven(t *testing.T) {
292296 case "upload-datadog-embedded-priority" :
293297 // Both embedded metadata and YAML file (embedded takes priority)
294298 // This tests that embedded metadata is prioritized over external YAML when both are present, proving precedence
295- dumpFilePath = generateMockTSDumpWithEmbeddedMetadata (t , d .Input )
299+ dumpFilePath = generateMockTSDumpFromCSV (t , d .Input , withEmbeddedMetadata () )
296300
297301 yamlContent := `1: "99"
2983022: "99"
@@ -326,53 +330,95 @@ func TestTSDumpUploadWithEmbeddedMetadataDataDriven(t *testing.T) {
326330 })
327331}
328332
329- // generateMockTSDumpWithEmbeddedMetadata creates a mock tsdump GOB file that includes
330- // embedded store-to-node mapping metadata along with time series data from CSV input.
331- func generateMockTSDumpWithEmbeddedMetadata (t * testing.T , csvInput string ) string {
332- t .Helper ()
333+ // mockTSDumpOption represents a functional option for configuring tsdump generation.
334+ type mockTSDumpOption func (* mockTSDumpConfig )
333335
334- tmpFile , err := os .CreateTemp ("" , "mock_tsdump_with_metadata_*.gob" )
335- require .NoError (t , err )
336- defer tmpFile .Close ()
336+ type mockTSDumpConfig struct {
337+ compression string
338+ metadata * tsdumpmeta.Metadata
339+ }
337340
338- metadata := tsdumpmeta.Metadata {
339- Version : "v23.1.0" ,
340- StoreToNodeMap : map [string ]string {
341- "1" : "1" ,
342- "2" : "1" ,
343- "3" : "2" ,
344- },
345- CreatedAt : timeutil .Unix (1609459200 , 0 ),
341+ // withCompression sets the compression format for the tsdump file.
342+ func withCompression (compression string ) mockTSDumpOption {
343+ return func (c * mockTSDumpConfig ) {
344+ c .compression = compression
346345 }
347- err = tsdumpmeta .Write (tmpFile , metadata )
348- require .NoError (t , err )
349-
350- encoder := gob .NewEncoder (tmpFile )
351- writeTimeSeriesDataFromCSV (t , csvInput , encoder )
346+ }
352347
353- t .Cleanup (func () {
354- require .NoError (t , os .Remove (tmpFile .Name ()), "failed to remove temporary file" )
355- })
356- return tmpFile .Name ()
348+ // withEmbeddedMetadata adds embedded metadata to the tsdump file.
349+ func withEmbeddedMetadata () mockTSDumpOption {
350+ return func (c * mockTSDumpConfig ) {
351+ c .metadata = & tsdumpmeta.Metadata {
352+ Version : "v23.1.0" ,
353+ StoreToNodeMap : map [string ]string {
354+ "1" : "1" ,
355+ "2" : "1" ,
356+ "3" : "2" ,
357+ },
358+ CreatedAt : timeutil .Unix (1609459200 , 0 ),
359+ }
360+ }
357361}
358362
359363// generateMockTSDumpFromCSV creates a mock tsdump file from CSV input string.
360364// CSV format: metric_name,timestamp,source,value
361365// Example: cr.node.admission.admitted.elastic-cpu,2025-05-26T08:32:00Z,1,1
362366// NOTE: this is the same format generated by the `cockroach tsdump` command
363367// when --format=csv is used.
364- func generateMockTSDumpFromCSV (t * testing.T , csvInput string ) string {
368+ // Options can be used to configure compression and embedded metadata.
369+ func generateMockTSDumpFromCSV (t * testing.T , csvInput string , options ... mockTSDumpOption ) string {
365370 t .Helper ()
366371
367- // Create temporary file
368- tmpFile , err := os .CreateTemp ("" , "mock_tsdump_*.gob" )
372+ // Apply options to config
373+ config := & mockTSDumpConfig {}
374+ for _ , option := range options {
375+ option (config )
376+ }
377+
378+ // Parse CSV data from input string
379+ reader := csv .NewReader (strings .NewReader (csvInput ))
380+ csvData , err := reader .ReadAll ()
381+ require .NoError (t , err )
382+ require .Greater (t , len (csvData ), 0 , "CSV input must have at least one data row" )
383+
384+ // Create file and encoder based on compression in single switch
385+ var filePattern string
386+ var encoder * gob.Encoder
387+ tmpFile , err := os .CreateTemp ("" , filePattern )
369388 require .NoError (t , err )
370389 defer tmpFile .Close ()
371390
372- // Create gob encoder
373- encoder := gob .NewEncoder (tmpFile )
374- writeTimeSeriesDataFromCSV (t , csvInput , encoder )
391+ if config .metadata != nil {
392+ err = tsdumpmeta .Write (tmpFile , * config .metadata )
393+ require .NoError (t , err )
394+ }
395+
396+ switch config .compression {
397+ case "gzip" :
398+ filePattern = "mock_tsdump_*.gob.gz"
399+ gzipWriter := gzip .NewWriter (tmpFile )
400+ defer func () {
401+ require .NoError (t , gzipWriter .Close (), "failed to close gzip writer" )
402+ }()
403+ encoder = gob .NewEncoder (gzipWriter )
404+
405+ case "zip" :
406+ filePattern = "mock_tsdump_*.gob.zip"
407+ zipWriter := zip .NewWriter (tmpFile )
408+ defer func () {
409+ require .NoError (t , zipWriter .Close (), "failed to close zip writer" )
410+ }()
411+
412+ writer , err := zipWriter .Create ("tsdump.gob" )
413+ require .NoError (t , err )
414+ encoder = gob .NewEncoder (writer )
415+
416+ default :
417+ filePattern = "mock_tsdump_*.gob"
418+ encoder = gob .NewEncoder (tmpFile )
419+ }
375420
421+ writeTimeSeriesDataFromCSV (t , csvInput , encoder )
376422 t .Cleanup (func () {
377423 require .NoError (t , os .Remove (tmpFile .Name ()), "failed to remove temporary file" )
378424 })
![craig[bot]](https://github.githubassets.com/images/gravatars/gravatar-user-420.png?size=40){kind=avatar}
0 commit comments