From 772de83a0631da97cb8c7bb25be5be11db505888 Mon Sep 17 00:00:00 2001
From: Lorenzo Primiterra <lorenzo@openobservatory.org>
Date: Fri, 26 Feb 2021 09:00:25 +0000
Subject: [PATCH] oonimkall: add wrappers for /test-list/urls (#221)

* First wrapper implementation

* Adding CountryCode parameter

* requested changes

* Update pkg/oonimkall/session.go

Co-authored-by: Simone Basso <bassosimone@gmail.com>

* Requested changes

* Update pkg/oonimkall/session_integration_test.go

Co-authored-by: Simone Basso <bassosimone@gmail.com>

* Update pkg/oonimkall/session_integration_test.go

Co-authored-by: Simone Basso <bassosimone@gmail.com>

* Remove duplicated code

* Apply suggestions from code review

Co-authored-by: Simone Basso <bassosimone@gmail.com>
---
 pkg/oonimkall/session.go                  | 67 +++++++++++++++++++++++
 pkg/oonimkall/session_integration_test.go | 54 ++++++++++++++++++
 2 files changed, 121 insertions(+)

diff --git a/pkg/oonimkall/session.go b/pkg/oonimkall/session.go
index 4a58dd1..ad47fc8 100644
--- a/pkg/oonimkall/session.go
+++ b/pkg/oonimkall/session.go
@@ -375,3 +375,70 @@ func (sess *Session) CheckIn(ctx *Context, config *CheckInConfig) (*CheckInInfo,
 		WebConnectivity: newCheckInInfoWebConnectivity(result.WebConnectivity),
 	}, nil
 }
+
+// URLListConfig contains configuration for fetching the URL list.
+type URLListConfig struct {
+	Categories  []string // Categories to query for (empty means all)
+	CountryCode string   // CountryCode is the optional country code
+	Limit       int64    // Max number of URLs (<= 0 means no limit)
+}
+
+// URLListResult contains the URLs returned from the FetchURL API
+type URLListResult struct {
+	Results []model.URLInfo
+}
+
+// AddCategory adds category code to the array in URLListConfig
+func (ckw *URLListConfig) AddCategory(cat string) {
+	ckw.Categories = append(ckw.Categories, cat)
+}
+
+// At gets the URLInfo at position idx from CheckInInfoWebConnectivity.URLs. It returns
+// nil if you are using an outs of bound index.
+func (ckw *URLListResult) At(idx int64) *URLInfo {
+	if idx < 0 || int(idx) >= len(ckw.Results) {
+		return nil
+	}
+	w := ckw.Results[idx]
+	return &URLInfo{
+		CategoryCode: w.CategoryCode,
+		CountryCode:  w.CountryCode,
+		URL:          w.URL,
+	}
+}
+
+// Size returns the number of URLs.
+func (ckw *URLListResult) Size() int64 {
+	return int64(len(ckw.Results))
+}
+
+// FetchURLList fetches the list of URLs to test
+func (sess *Session) FetchURLList(ctx *Context, config *URLListConfig) (*URLListResult, error) {
+	sess.mtx.Lock()
+	defer sess.mtx.Unlock()
+	psc, err := sess.sessp.NewProbeServicesClient(ctx.ctx)
+	if err != nil {
+		return nil, err
+	}
+	if config.CountryCode == "" {
+		config.CountryCode = "XX"
+		info, err := sess.sessp.LookupLocationContext(ctx.ctx)
+		if err == nil && info != nil {
+			config.CountryCode = info.CountryCode
+		}
+	}
+
+	cfg := model.URLListConfig{
+		Categories:  config.Categories,
+		CountryCode: config.CountryCode,
+		Limit:       config.Limit,
+	}
+
+	result, err := psc.FetchURLList(ctx.ctx, cfg)
+	if err != nil {
+		return nil, err
+	}
+	return &URLListResult{
+		Results: result,
+	}, nil
+}
diff --git a/pkg/oonimkall/session_integration_test.go b/pkg/oonimkall/session_integration_test.go
index 8cd8750..967243d 100644
--- a/pkg/oonimkall/session_integration_test.go
+++ b/pkg/oonimkall/session_integration_test.go
@@ -419,6 +419,60 @@ func TestCheckInNoParams(t *testing.T) {
 	}
 }
 
+func TestFetchURLListSuccess(t *testing.T) {
+	sess, err := NewSession()
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := sess.NewContext()
+	config := oonimkall.URLListConfig{
+		Limit: 10,
+	}
+	config.AddCategory("NEWS")
+	config.AddCategory("CULTR")
+	result, err := sess.FetchURLList(ctx, &config)
+	if err != nil {
+		t.Fatalf("unexpected error: %+v", err)
+	}
+	if result == nil || result.Results == nil {
+		t.Fatal("got nil result")
+	}
+	for _, entry := range result.Results {
+		if entry.CategoryCode != "NEWS" && entry.CategoryCode != "CULTR" {
+			t.Fatalf("unexpected category code: %+v", entry)
+		}
+	}
+}
+
+func TestFetchURLListWithCC(t *testing.T) {
+	sess, err := NewSession()
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := sess.NewContext()
+	config := oonimkall.URLListConfig{
+		CountryCode: "IT",
+	}
+	config.AddCategory("NEWS")
+	config.AddCategory("CULTR")
+	result, err := sess.FetchURLList(ctx, &config)
+	if err != nil {
+		t.Fatalf("unexpected error: %+v", err)
+	}
+	if result == nil || result.Results == nil {
+		t.Fatal("got nil result")
+	}
+	found := false
+	for _, entry := range result.Results {
+		if entry.CountryCode == "IT" {
+			found = true
+		}
+	}
+	if !found {
+		t.Fatalf("not found url for country code: IT")
+	}
+}
+
 func TestMain(m *testing.M) {
 	// Here we're basically testing whether eventually the finalizers
 	// will run and the number of active sessions and cancels will become