Add support for passing the limit and category code to the testlists pkg

orchestra/testlists/testlists.go

@@ -3,6 +3,7 @@ package testlists
 
 import (
 	"context"
+	"strings"
 	"net/http"
 	"net/url"
 
@@ -45,6 +46,9 @@ type Client struct {
 
 	// UserAgent is the user agent to use.
 	UserAgent string
+
+	// EnabledCategories is a list of category codes that are enabled
+	EnabledCategories []string
 }
 
 // NewClient creates a new client in the context of the given session.
@@ -57,15 +61,26 @@ func NewClient(sess *session.Session) *Client {
 	}
 }
 
+// SetEnabledCategories configures the client category codes
+func (c *Client) SetEnabledCategories(categories []string) error {
+	c.EnabledCategories = categories
+}
+
 // Do retrieves the test list for the specified country.
 func (c *Client) Do(
-	ctx context.Context, countryCode string,
+	ctx context.Context, countryCode string, limit int,
 ) ([]URLInfo, error) {
 	var resp response
 	query := url.Values{}
 	if countryCode != "" {
 		query.Set("probe_cc", countryCode)
 	}
+	if (limit > 0) {
+		query.Set("limit", limit)
+	}
+	if (len(c.EnabledCategories) > 0) {
+		query.Set("category_codes", strings.Join(c.EnabledCategories, ","))
+	}
 	err := (&jsonapi.Client{
 		BaseURL:    c.BaseURL,
 		HTTPClient: c.HTTPClient,

orchestra/testlists/testlists_test.go

@@ -20,7 +20,7 @@ func makeClient() *testlists.Client {
 
 func TestIntegration(t *testing.T) {
 	log.SetLevel(log.DebugLevel)
-	urls, err := makeClient().Do(context.Background(), "IT")
+	urls, err := makeClient().Do(context.Background(), "IT", 0)
 	if err != nil {
 		t.Fatal(err)
 	}