Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make test utils exportable #289

Merged
merged 13 commits into from
Dec 3, 2024
12 changes: 6 additions & 6 deletions cmd/command_utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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 {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand All @@ -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:
//
Expand All @@ -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)
Expand Down
12 changes: 6 additions & 6 deletions cmd/export_assets.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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)
}
},
}
Expand Down
36 changes: 18 additions & 18 deletions cmd/export_assets_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
12 changes: 6 additions & 6 deletions cmd/export_contract_events.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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
Expand All @@ -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)
}

},
Expand Down
12 changes: 6 additions & 6 deletions cmd/export_contract_events_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
12 changes: 6 additions & 6 deletions cmd/export_effects.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
}
},
}
Expand Down
36 changes: 18 additions & 18 deletions cmd/export_effects_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
16 changes: 8 additions & 8 deletions cmd/export_ledger_entry_changes.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
}
}

Expand Down
Loading
Loading