Make test utils exportable (#289)

* export cmd utils * Move common functions to command utils * export more funcs * Export more funcs * update * export struct vars * update * make executable name arg * add arg for using parent dir * log * run command directly * Remove debug statements * Separate util file for test utils
stellar · Dec 3, 2024 · 53929d8 · 53929d8
1 parent 204d343
commit 53929d8
Show file tree

Hide file tree

Showing 23 changed files with 477 additions and 462 deletions.
diff --git a/cmd/command_utils.go b/cmd/command_utils.go
@@ -28,7 +28,7 @@ func createOutputFile(filepath string) error {
 	return nil
 }
 
-func mustOutFile(path string) *os.File {
+func MustOutFile(path string) *os.File {
 	absolutePath, err := filepath.Abs(path)
 	if err != nil {
 		cmdLogger.Fatal("could not get absolute filepath: ", err)
@@ -52,7 +52,7 @@ func mustOutFile(path string) *os.File {
 	return outFile
 }
 
-func exportEntry(entry interface{}, outFile *os.File, extra map[string]string) (int, error) {
+func ExportEntry(entry interface{}, outFile *os.File, extra map[string]string) (int, error) {
 	// This extra marshalling/unmarshalling is silly, but it's required to properly handle the null.[String|Int*] types, and add the extra fields.
 	m, err := json.Marshal(entry)
 	if err != nil {
@@ -87,7 +87,7 @@ func exportEntry(entry interface{}, outFile *os.File, extra map[string]string) (
 }
 
 // Prints the number of attempted, failed, and successful transformations as a JSON object
-func printTransformStats(attempts, failures int) {
+func PrintTransformStats(attempts, failures int) {
 	resultsMap := map[string]int{
 		"attempted_transforms":  attempts,
 		"failed_transforms":     failures,
@@ -120,7 +120,7 @@ func deleteLocalFiles(path string) error {
 	return nil
 }
 
-func maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path string) {
+func MaybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path string) {
 	if cloudProvider == "" {
 		cmdLogger.Info("No cloud provider specified for upload. Skipping upload.")
 		return
@@ -145,7 +145,7 @@ func maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path strin
 	}
 }
 
-// writeParquet creates the parquet file and writes the exported data into it.
+// WriteParquet creates the parquet file and writes the exported data into it.
 //
 // Parameters:
 //
@@ -159,7 +159,7 @@ func maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path strin
 //	Errors:
 //
 //	stellar-etl will log a Fatal error and stop in the case it cannot create or write to the parquet file
-func writeParquet(data []transform.SchemaParquet, path string, schema interface{}) {
+func WriteParquet(data []transform.SchemaParquet, path string, schema interface{}) {
 	parquetFile, err := local.NewLocalFileWriter(path)
 	if err != nil {
 		cmdLogger.Fatal("could not create parquet file: ", err)

diff --git a/cmd/export_assets.go b/cmd/export_assets.go
@@ -22,7 +22,7 @@ var assetsCmd = &cobra.Command{
 		cloudStorageBucket, cloudCredentials, cloudProvider := utils.MustCloudStorageFlags(cmd.Flags(), cmdLogger)
 		env := utils.GetEnvironmentDetails(commonArgs)
 
-		outFile := mustOutFile(path)
+		outFile := MustOutFile(path)
 
 		var paymentOps []input.AssetTransformInput
 		var err error
@@ -56,7 +56,7 @@ var assetsCmd = &cobra.Command{
 			}
 
 			seenIDs[transformed.AssetID] = true
-			numBytes, err := exportEntry(transformed, outFile, commonArgs.Extra)
+			numBytes, err := ExportEntry(transformed, outFile, commonArgs.Extra)
 			if err != nil {
 				cmdLogger.LogError(err)
 				numFailures += 1
@@ -72,13 +72,13 @@ var assetsCmd = &cobra.Command{
 		outFile.Close()
 		cmdLogger.Infof("%d bytes written to %s", totalNumBytes, outFile.Name())
 
-		printTransformStats(len(paymentOps), numFailures)
+		PrintTransformStats(len(paymentOps), numFailures)
 
-		maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path)
+		MaybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path)
 
 		if commonArgs.WriteParquet {
-			writeParquet(transformedAssets, parquetPath, new(transform.AssetOutputParquet))
-			maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, parquetPath)
+			WriteParquet(transformedAssets, parquetPath, new(transform.AssetOutputParquet))
+			MaybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, parquetPath)
 		}
 	},
 }

diff --git a/cmd/export_assets_test.go b/cmd/export_assets_test.go
@@ -5,34 +5,34 @@ import (
 )
 
 func TestExportAssets(t *testing.T) {
-	tests := []cliTest{
+	tests := []CliTest{
 		{
-			name:    "assets from one ledger",
-			args:    []string{"export_assets", "-s", "30820015", "-e", "30820015", "-o", gotTestDir(t, "one_ledger_assets.txt")},
-			golden:  "one_ledger_assets.golden",
-			wantErr: nil,
+			Name:    "assets from one ledger",
+			Args:    []string{"export_assets", "-s", "30820015", "-e", "30820015", "-o", GotTestDir(t, "one_ledger_assets.txt")},
+			Golden:  "one_ledger_assets.golden",
+			WantErr: nil,
 		},
 		{
-			name:    "assets from 10 ledgers",
-			args:    []string{"export_assets", "-s", "30822015", "-e", "30822025", "-o", gotTestDir(t, "10_ledgers_assets.txt")},
-			golden:  "10_ledgers_assets.golden",
-			wantErr: nil,
+			Name:    "assets from 10 ledgers",
+			Args:    []string{"export_assets", "-s", "30822015", "-e", "30822025", "-o", GotTestDir(t, "10_ledgers_assets.txt")},
+			Golden:  "10_ledgers_assets.golden",
+			WantErr: nil,
 		},
 		{
-			name:    "range too large",
-			args:    []string{"export_assets", "-s", "30822015", "-e", "30822025", "-l", "5", "-o", gotTestDir(t, "large_range_assets.txt")},
-			golden:  "large_range_assets.golden",
-			wantErr: nil,
+			Name:    "range too large",
+			Args:    []string{"export_assets", "-s", "30822015", "-e", "30822025", "-l", "5", "-o", GotTestDir(t, "large_range_assets.txt")},
+			Golden:  "large_range_assets.golden",
+			WantErr: nil,
 		},
 		{
-			name:    "ledger with no assets",
-			args:    []string{"export_assets", "-s", "10363513", "-e", "10363513", "-o", gotTestDir(t, "ledger_no_assets.txt")},
-			golden:  "ledger_no_assets.golden",
-			wantErr: nil,
+			Name:    "ledger with no assets",
+			Args:    []string{"export_assets", "-s", "10363513", "-e", "10363513", "-o", GotTestDir(t, "ledger_no_assets.txt")},
+			Golden:  "ledger_no_assets.golden",
+			WantErr: nil,
 		},
 	}
 
 	for _, test := range tests {
-		runCLITest(t, test, "testdata/assets/")
+		RunCLITest(t, test, "testdata/assets/", "", false)
 	}
 }
diff --git a/cmd/export_contract_events.go b/cmd/export_contract_events.go
@@ -27,7 +27,7 @@ var contractEventsCmd = &cobra.Command{
 			cmdLogger.Fatal("could not read transactions: ", err)
 		}
 
-		outFile := mustOutFile(cmdArgs.Path)
+		outFile := MustOutFile(cmdArgs.Path)
 		numFailures := 0
 		var transformedEvents []transform.SchemaParquet
 		for _, transformInput := range transactions {
@@ -40,7 +40,7 @@ var contractEventsCmd = &cobra.Command{
 			}
 
 			for _, contractEvent := range transformed {
-				_, err := exportEntry(contractEvent, outFile, cmdArgs.Extra)
+				_, err := ExportEntry(contractEvent, outFile, cmdArgs.Extra)
 				if err != nil {
 					cmdLogger.LogError(fmt.Errorf("could not export contract event: %v", err))
 					numFailures += 1
@@ -56,13 +56,13 @@ var contractEventsCmd = &cobra.Command{
 
 		outFile.Close()
 
-		printTransformStats(len(transactions), numFailures)
+		PrintTransformStats(len(transactions), numFailures)
 
-		maybeUpload(cmdArgs.Credentials, cmdArgs.Bucket, cmdArgs.Provider, cmdArgs.Path)
+		MaybeUpload(cmdArgs.Credentials, cmdArgs.Bucket, cmdArgs.Provider, cmdArgs.Path)
 
 		if commonArgs.WriteParquet {
-			writeParquet(transformedEvents, cmdArgs.ParquetPath, new(transform.ContractEventOutputParquet))
-			maybeUpload(cmdArgs.Credentials, cmdArgs.Bucket, cmdArgs.Provider, cmdArgs.ParquetPath)
+			WriteParquet(transformedEvents, cmdArgs.ParquetPath, new(transform.ContractEventOutputParquet))
+			MaybeUpload(cmdArgs.Credentials, cmdArgs.Bucket, cmdArgs.Provider, cmdArgs.ParquetPath)
 		}
 
 	},

diff --git a/cmd/export_contract_events_test.go b/cmd/export_contract_events_test.go
@@ -5,16 +5,16 @@ import (
 )
 
 func TestExportContractEvents(t *testing.T) {
-	tests := []cliTest{
+	tests := []CliTest{
 		{
-			name:    "contract events from multiple ledger",
-			args:    []string{"export_contract_events", "-s", "52271338", "-e", "52271350", "-o", gotTestDir(t, "large_range_ledger_txs.txt")},
-			golden:  "large_range_ledger_txs.golden",
-			wantErr: nil,
+			Name:    "contract events from multiple ledger",
+			Args:    []string{"export_contract_events", "-s", "52271338", "-e", "52271350", "-o", GotTestDir(t, "large_range_ledger_txs.txt")},
+			Golden:  "large_range_ledger_txs.golden",
+			WantErr: nil,
 		},
 	}
 
 	for _, test := range tests {
-		runCLITest(t, test, "testdata/contract_events/")
+		RunCLITest(t, test, "testdata/contract_events/", "", false)
 	}
 }
diff --git a/cmd/export_effects.go b/cmd/export_effects.go
@@ -27,7 +27,7 @@ var effectsCmd = &cobra.Command{
 			cmdLogger.Fatalf("could not read transactions in [%d, %d] (limit=%d): %v", startNum, commonArgs.EndNum, limit, err)
 		}
 
-		outFile := mustOutFile(path)
+		outFile := MustOutFile(path)
 		numFailures := 0
 		totalNumBytes := 0
 		var transformedEffects []transform.SchemaParquet
@@ -42,7 +42,7 @@ var effectsCmd = &cobra.Command{
 			}
 
 			for _, transformed := range effects {
-				numBytes, err := exportEntry(transformed, outFile, commonArgs.Extra)
+				numBytes, err := ExportEntry(transformed, outFile, commonArgs.Extra)
 				if err != nil {
 					cmdLogger.LogError(err)
 					numFailures += 1
@@ -59,13 +59,13 @@ var effectsCmd = &cobra.Command{
 		outFile.Close()
 		cmdLogger.Info("Number of bytes written: ", totalNumBytes)
 
-		printTransformStats(len(transactions), numFailures)
+		PrintTransformStats(len(transactions), numFailures)
 
-		maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path)
+		MaybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path)
 
 		if commonArgs.WriteParquet {
-			writeParquet(transformedEffects, parquetPath, new(transform.EffectOutputParquet))
-			maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, parquetPath)
+			WriteParquet(transformedEffects, parquetPath, new(transform.EffectOutputParquet))
+			MaybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, parquetPath)
 		}
 	},
 }

diff --git a/cmd/export_effects_test.go b/cmd/export_effects_test.go
@@ -5,34 +5,34 @@ import (
 )
 
 func TestExportEffects(t *testing.T) {
-	tests := []cliTest{
+	tests := []CliTest{
 		{
-			name:    "effects from one ledger",
-			args:    []string{"export_effects", "-s", "30820015", "-e", "30820015", "-o", gotTestDir(t, "one_ledger_effects.txt")},
-			golden:  "one_ledger_effects.golden",
-			wantErr: nil,
+			Name:    "effects from one ledger",
+			Args:    []string{"export_effects", "-s", "30820015", "-e", "30820015", "-o", GotTestDir(t, "one_ledger_effects.txt")},
+			Golden:  "one_ledger_effects.golden",
+			WantErr: nil,
 		},
 		{
-			name:    "effects from 10 ledgers",
-			args:    []string{"export_effects", "-s", "30822015", "-e", "30822025", "-o", gotTestDir(t, "10_ledgers_effects.txt")},
-			golden:  "10_ledgers_effects.golden",
-			wantErr: nil,
+			Name:    "effects from 10 ledgers",
+			Args:    []string{"export_effects", "-s", "30822015", "-e", "30822025", "-o", GotTestDir(t, "10_ledgers_effects.txt")},
+			Golden:  "10_ledgers_effects.golden",
+			WantErr: nil,
 		},
 		{
-			name:    "range too large",
-			args:    []string{"export_effects", "-s", "25820678", "-e", "25821678", "-l", "5", "-o", gotTestDir(t, "large_range_effects.txt")},
-			golden:  "large_range_effects.golden",
-			wantErr: nil,
+			Name:    "range too large",
+			Args:    []string{"export_effects", "-s", "25820678", "-e", "25821678", "-l", "5", "-o", GotTestDir(t, "large_range_effects.txt")},
+			Golden:  "large_range_effects.golden",
+			WantErr: nil,
 		},
 		{
-			name:    "ledger with no effects",
-			args:    []string{"export_effects", "-s", "10363513", "-e", "10363513", "-o", gotTestDir(t, "ledger_no_effects.txt")},
-			golden:  "ledger_no_effects.golden",
-			wantErr: nil,
+			Name:    "ledger with no effects",
+			Args:    []string{"export_effects", "-s", "10363513", "-e", "10363513", "-o", GotTestDir(t, "ledger_no_effects.txt")},
+			Golden:  "ledger_no_effects.golden",
+			WantErr: nil,
 		},
 	}
 
 	for _, test := range tests {
-		runCLITest(t, test, "testdata/effects/")
+		RunCLITest(t, test, "testdata/effects/", "", false)
 	}
 }
diff --git a/cmd/export_ledger_entry_changes.go b/cmd/export_ledger_entry_changes.go
@@ -287,25 +287,25 @@ func exportTransformedData(
 	transformedOutput map[string][]interface{},
 	cloudCredentials, cloudStorageBucket, cloudProvider string,
 	extra map[string]string,
-	WriteParquet bool) error {
+	writeParquet bool) error {
 
 	for resource, output := range transformedOutput {
 		// Filenames are typically exclusive of end point. This processor
 		// is different and we have to increment by 1 since the end batch number
 		// is included in this filename.
 		path := filepath.Join(folderPath, exportFilename(start, end+1, resource))
 		parquetPath := filepath.Join(parquetFolderPath, exportParquetFilename(start, end+1, resource))
-		outFile := mustOutFile(path)
+		outFile := MustOutFile(path)
 		var transformedResource []transform.SchemaParquet
 		var parquetSchema interface{}
 		var skip bool
 		for _, o := range output {
-			_, err := exportEntry(o, outFile, extra)
+			_, err := ExportEntry(o, outFile, extra)
 			if err != nil {
 				return err
 			}
 
-			if WriteParquet {
+			if writeParquet {
 				switch v := o.(type) {
 				case transform.AccountOutput:
 					transformedResource = append(transformedResource, v)
@@ -352,11 +352,11 @@ func exportTransformedData(
 			}
 		}
 
-		maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path)
+		MaybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, path)
 
-		if !skip && WriteParquet {
-			writeParquet(transformedResource, parquetPath, parquetSchema)
-			maybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, parquetPath)
+		if !skip && writeParquet {
+			WriteParquet(transformedResource, parquetPath, parquetSchema)
+			MaybeUpload(cloudCredentials, cloudStorageBucket, cloudProvider, parquetPath)
 		}
 	}