github.com/mweagle/Sparta@v1.15.0/archetype/firehose_test.go (about)

     1  package archetype
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"io/ioutil"
     7  	"strings"
     8  	"testing"
     9  
    10  	awsEvents "github.com/aws/aws-lambda-go/events"
    11  	awsEventsTest "github.com/aws/aws-lambda-go/events/test"
    12  	"github.com/pkg/errors"
    13  )
    14  
    15  var firehoseTests = []struct {
    16  	sourceInputPath string
    17  	templatePath    string
    18  	predicate       testPredicate
    19  }{
    20  	{
    21  		"test/records.json",
    22  		"test/testdata.transform",
    23  		okPredicate,
    24  	},
    25  	{
    26  		"test/records-sm.json",
    27  		"test/drop.transform",
    28  		verifyPredicate("Dropped", 3),
    29  	},
    30  	{
    31  		"test/records.json",
    32  		"test/conditional.transform",
    33  		okPredicate,
    34  	},
    35  	{
    36  		"test/records-sm.json",
    37  		"test/conditional.transform",
    38  		verifyPredicate("Dropped", 1),
    39  	},
    40  	{
    41  		"test/records-sm.json",
    42  		"test/regexp.transform",
    43  		verifyPredicate("sector", 3),
    44  	},
    45  	{
    46  		"test/records-sm.json",
    47  		"test/regexp.transform",
    48  		verifyPredicate("FINANCIAL", 2),
    49  	},
    50  	{
    51  		"test/records-sm.json",
    52  		"test/regexpEmpty.transform",
    53  		verifyPredicate("TECHNOLOGY", 0),
    54  	},
    55  }
    56  
    57  type testPredicate func(t *testing.T, response *awsEvents.KinesisFirehoseResponse) error
    58  
    59  func okPredicate(t *testing.T, response *awsEvents.KinesisFirehoseResponse) error {
    60  
    61  	for _, eachEntry := range response.Records {
    62  		jsonMap := make(map[string]interface{})
    63  		unmarshalErr := json.Unmarshal(eachEntry.Data, &jsonMap)
    64  		if unmarshalErr != nil {
    65  			return unmarshalErr
    66  		}
    67  		//t.Logf("Record: %#v\n", jsonMap)
    68  	}
    69  	return nil
    70  }
    71  
    72  func verifyPredicate(value string, expectedCount int) testPredicate {
    73  	return func(t *testing.T, response *awsEvents.KinesisFirehoseResponse) error {
    74  		counter := 0
    75  		jsonMap := make(map[string]interface{})
    76  		for _, eachEntry := range response.Records {
    77  			rawData := string(eachEntry.Data)
    78  
    79  			unmarshalErr := json.Unmarshal(eachEntry.Data, &jsonMap)
    80  			if unmarshalErr == nil {
    81  				jsonData, _ := json.MarshalIndent(jsonMap, "", " ")
    82  				rawData = string(jsonData)
    83  			}
    84  			//t.Logf("RAW DATA: %s\n", string(rawData))
    85  			counter += strings.Count(rawData, value)
    86  		}
    87  		jsonData, _ := json.MarshalIndent(response, "", " ")
    88  		counter += strings.Count(string(jsonData), value)
    89  		// t.Log(string(jsonData))
    90  		// counter += strings.Count(string(jsonData), value)
    91  
    92  		if counter != expectedCount {
    93  			return errors.Errorf("Invalid count. Expected: %d, Found: %d for value: %s",
    94  				expectedCount,
    95  				counter,
    96  				value)
    97  		}
    98  		return nil
    99  	}
   100  }
   101  
   102  func testData(t *testing.T, inputPath string) *awsEvents.KinesisFirehoseEvent {
   103  
   104  	// 1. read JSON from file
   105  	inputJSON := awsEventsTest.ReadJSONFromFile(t, inputPath)
   106  
   107  	// 2. de-serialize into Go object
   108  	var inputEvent awsEvents.KinesisFirehoseEvent
   109  	if err := json.Unmarshal(inputJSON, &inputEvent); err != nil {
   110  		t.Fatalf("could not unmarshal event. details: %v", err)
   111  		return nil
   112  	}
   113  	return &inputEvent
   114  }
   115  
   116  func testTemplate(t *testing.T, xformPath string) []byte {
   117  	templateBytes, templateBytesErr := ioutil.ReadFile(xformPath)
   118  	if templateBytesErr != nil {
   119  		t.Fatalf("Could not read template data. Error: %v", templateBytesErr)
   120  		return nil
   121  	}
   122  	return templateBytes
   123  }
   124  
   125  func TestTransforms(t *testing.T) {
   126  
   127  	for _, tt := range firehoseTests {
   128  		t.Run(tt.templatePath, func(t *testing.T) {
   129  			data := testData(t, tt.sourceInputPath)
   130  			xformTemplate := testTemplate(t, tt.templatePath)
   131  			ctx := context.Background()
   132  			response, responseErr := ApplyTransformToKinesisFirehoseEvent(ctx,
   133  				xformTemplate,
   134  				*data)
   135  			if responseErr != nil {
   136  				t.Fatal(responseErr)
   137  				return
   138  			}
   139  			// Unmarshal everything...
   140  			predicateErr := tt.predicate(t, response)
   141  			if predicateErr != nil {
   142  				t.Fatal(predicateErr)
   143  				return
   144  			}
   145  		})
   146  	}
   147  }
   148  
   149  func TestLambdaTransform(t *testing.T) {
   150  	lambdaTransform := func(ctx context.Context,
   151  		kinesisRecord *awsEvents.KinesisFirehoseEventRecord) (*awsEvents.KinesisFirehoseResponseRecord, error) {
   152  
   153  		return &awsEvents.KinesisFirehoseResponseRecord{
   154  			RecordID: kinesisRecord.RecordID,
   155  			Result:   awsEvents.KinesisFirehoseTransformedStateOk,
   156  			Data:     kinesisRecord.Data,
   157  		}, nil
   158  	}
   159  	ctx := context.Background()
   160  
   161  	for _, tt := range firehoseTests {
   162  		t.Run(tt.templatePath, func(t *testing.T) {
   163  			data := testData(t, tt.sourceInputPath)
   164  
   165  			for _, eachRecord := range data.Records {
   166  				xformed, xformedErr := lambdaTransform(ctx, &eachRecord)
   167  				if xformedErr != nil {
   168  					t.Fatalf("Failed to transform")
   169  				}
   170  				if xformed.Result != awsEvents.KinesisFirehoseTransformedStateOk {
   171  					t.Fatalf("Failed to successful process record")
   172  				}
   173  			}
   174  		})
   175  	}
   176  
   177  }