POC to production - transition of code

At Banzai Cloud we’re always open to experimenting with and integrating new software (tools, products). We also love to validate our new ideas by quickly implementing “proof of concept” projects. Even though we used five or so programming languages while building the Pipeline Platform, we love and use Golang the most. While these PoC projects are not intended for production use, they often serve as the basis for it. When this is the case, the PoC code needs to be refactored - or prepared for production. The focus changes from quick trials and idea validation to reliability, scalability, extensibility, testability and maintainability.

This post serves as a brief examination of the process by which a small part of one of our components transitioned from PoC to production - with an emphasis on (unit) testability (i.e. we want to make sure the component does what it’s supposed to).

Making code production-ready is driven by the paradigm that code of good quality can be easily unit tested and, vice versa, that code that is easily unit tested is of good quality (in other words, tests are good).

A lot of material has been written about what unit tests are (for example: this), and about the best methods and tricks to improve unit test writing. Consequentially, this post won’t focus on these, but rather it will describe how a small part of our code became well testable.

The component we’ll be using as an example was in charge of retrieving information from AWS through the AWS Price List Service API, and of processing the results. The goal was to test the component’s functionality in isolation, without connecting it to other components.

The example below comes from our Kubernetes cluster infrastructure recommendation project, Telescopes. We’d spent some time playing with a PoC to recommend infrastructures for our Kubernetes node pools and that project eventually evolved into a key component of the Pipeline Platform, thus we needed to refactor it for better testability.

The initial state 🔗︎

One of the methods that needed to be tested looked like this:

// AppMethod builds the pricing reference, retrieves pricing information  and transforms the response
// this method is being refactored so that it can be tested
func (au *AppStruct) AppMethod() (*DomainStruct, error) {

	// start boilerplate
	s, err := session.NewSession(aws.NewConfig())
	if err != nil {
		return nil, fmt.Errorf("could not create session")
	}

	pr := pricing.New(s)
	// end  boilerplate!

	// method specific logic
	resp, err := pr.GetAttributeValues(&pricing.GetAttributeValuesInput{
		ServiceCode:   aws.String("AmazonEC2"),
		AttributeName: aws.String("memory"),
	})

	ds, err := au.transform(resp)
	if err != nil {
		return nil, fmt.Errorf("could not transform")
	}

	return ds, nil
}

Although not extremely complicated, there are several problems with this implementation and reasons why this method can’t be properly unit tested:

• First, it does too much; it builds the connection to AWS, retrieves the data, and performs logic.
• Second, it cannot imitate the AWS client, so this method will always try to connect to the “real” AWS - this makes it impossible to test the proprietary logic (transformation) in isolation.

Based on the above observations, let’s simplify the method body: extract the boilerplate and provide the method with the service itself:

 func (a *AppStruct) AppMethod(pr *pricing.Pricing) (*DomainStruct, error) {

 	resp, err := pr.GetAttributeValues(&pricing.GetAttributeValuesInput{
 		ServiceCode:   aws.String("AmazonEC2"),
 		AttributeName: aws.String("memory"),
 	})

 	ds, err := a.transform(resp)
    if err != nil {
        return nil, fmt.Errorf("could not transform")
    }

    return ds, nil
 }

 

The method body is more focused on its exact task - getting the data and processing it. All the boilerplate and logic that does not strictly belong to this method has been extracted. More importantly, the reference *pricing.Pricing is easily inferred. There are a few more ways to clearly compartmentalize responsibilities:

• by making the reference *pricing.Pricing into a field in AppStruct, so that it can be reused many times (in this case, the service instance can be initialized when creating the AppStruct instance)
• and by only providing the response to the method being discussed

Let’s stick with our first option; the code now looks like this:

type AppStruct struct {
	pricing *pricing.Pricing
}

....


func (a *AppStruct) AppMethod() (*DomainStruct, error) {

	resp, err := a.pricing.GetAttributeValues(&pricing.GetAttributeValuesInput{
		ServiceCode:   aws.String("AmazonEC2"),
		AttributeName: aws.String("memory"),
	})

	ds, err := a.transform(resp)
	if err != nil {
		return nil, fmt.Errorf("could not transform: %s", err.Error())
	}

	return ds, nil
}

Looking at the code, it’s clear that the pricing service collaborates with the method; this means we’ve found one of the component’s boundaries.

Now let’s see how this pricing service can be “mocked” and inferred into AppStruct. We need to infer a reference for pricing struct that is entirely under our control into AppStruct. This will let us define an internal interface that lists the operation on the pricing service:

type pricingSource interface {
	GetAttributeValues(input *pricing.GetAttributeValuesInput) (*pricing.GetAttributeValuesOutput, error)
}

Note that the original *pring.Pricing struct implements this interface

Let’s change the type of field in the AppStruct:

type AppStruct struct {
	pricing pricingSource
}

With this step, we almost cut off the method’s direct dependency on the aws-sdk-go library. (It’s also possible for us to remove dependencies on other lib-dependent objects like method arguments and the response type of the interface method; but because these objects can be easily mocked and due to the unnecessary transformations involved, we’ll stop here)

The code now looks like this:

type AppStruct struct {
	pricingSource pricingSource
}

type pricingSource interface {
	GetAttributeValues(input *pricing.GetAttributeValuesInput) (*pricing.GetAttributeValuesOutput, error)
}

func (a *AppStruct) AppMethod() (*DomainStruct, error) {

	resp, err := a.pricingSource.GetAttributeValues(&pricing.GetAttributeValuesInput{
		ServiceCode:   aws.String("AmazonEC2"),
		AttributeName: aws.String("memory"),
	})

	ds, err := a.transform(resp)
	if err != nil {
		return nil, fmt.Errorf("could not transform: %s", err.Error())
	}

	return ds, nil
}

The unit test for the above method:

// dummyPricing struct for mocking the pricing service
type dummyPricing struct {
	// add fields to drive the test scenarios if needed
}

// GetAttributeValues is the behaviour under our control; returns well controlled responses based on input (or other control flags)
func (d dummyPricing) GetAttributeValues(input *pricing.GetAttributeValuesInput) (*pricing.GetAttributeValuesOutput, error) {
	// return error / well defined pricing.GetAttributeValuesOutput reference as the test case needs it
	return nil, nil;
}

func TestAppStruct_AppMethod(t *testing.T) {
	tests := []struct {
		// insert the name of the test case
		name string
		// check and make sure the placeholder function is checking the results of the method call
		check func(domainStruct *DomainStruct, err error)
	}{
		{
			name: "testing",
			check: func(domainStruct *DomainStruct, err error) {
				// assert on the response received, eg:
				assert.Nil(t, err, "should not receive error")

			},
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			dps, err := NewAppStruct(dummyPricing{})
			if err != nil {
				t.Fatal("failed to initialize app: %s", err.Error())
			}
			// call the method under test
			ds, err := dps.AppMethod()

			// check the results
			test.check(ds, err)

		})
	}
}

The test above is a simple table test; any number of test cases can be defined to fully cover the method’s functionality. One trick we found useful was defining the checker method for each of the test cases, so assertions could be made accordingly.

To summarize, these are the steps of the code transition we’ve just described:

• identify the task the method is in charge of (ideally a method should perform one task)
• identify component boundaries / collaborators
• define (internal) interfaces that describe interactions between components (“declare” the boundaries)
• use interfaces in application code
• write tests that use a dummy implementation of the above interface
• refactor the application code (remove unreachable blocks, simplify the code, use appropriate blocks and constructs etc …)

Please find the sample code in our GitHub repository

Note: We are aware that there are lots of Go libraries that help out with mocking (and AWS does a particularly good job of this), and we deliberately tried to keep this example generic and simple. Happy testing.