ooni-probe-cli/internal/nettests/web_connectivity.go
Simone Basso a324e94973
refactor(webconnectivity): use engine.InputLoader (#179)
This diff pins to ooni/probe-engine@3049779878
and starts using the recently introduced probe-engine APIs.

Namely, here, we use the InputLoader for loading URLs.

I've confirmed manually everything is still working as intended.

Part of https://github.com/ooni/probe/issues/1283.

(In particular, the InputLoader is the abstraction allowing us to load
input from several sources, including command line flags and external
files.)
2020-11-30 15:39:25 +01:00

60 lines
1.6 KiB
Go

package nettests
import (
"context"
"github.com/apex/log"
"github.com/ooni/probe-cli/internal/database"
engine "github.com/ooni/probe-engine"
)
func lookupURLs(ctl *Controller, limit int64, categories []string) ([]string, map[int64]int64, error) {
inputloader := engine.NewInputLoader(engine.InputLoaderConfig{
InputPolicy: engine.InputRequired,
Session: ctl.Session,
URLCategories: categories,
URLLimit: limit,
})
testlist, err := inputloader.Load(context.Background())
var urls []string
urlIDMap := make(map[int64]int64)
if err != nil {
return nil, nil, err
}
for idx, url := range testlist {
log.Debugf("Going over URL %d", idx)
urlID, err := database.CreateOrUpdateURL(
ctl.Probe.DB(), url.URL, url.CategoryCode, url.CountryCode,
)
if err != nil {
log.Error("failed to add to the URL table")
return nil, nil, err
}
log.Debugf("Mapped URL %s to idx %d and urlID %d", url.URL, idx, urlID)
urlIDMap[int64(idx)] = urlID
urls = append(urls, url.URL)
}
return urls, urlIDMap, nil
}
// WebConnectivity test implementation
type WebConnectivity struct {
}
// Run starts the test
func (n WebConnectivity) Run(ctl *Controller) error {
log.Debugf("Enabled category codes are the following %v", ctl.Probe.Config().Nettests.WebsitesEnabledCategoryCodes)
urls, urlIDMap, err := lookupURLs(ctl, ctl.Probe.Config().Nettests.WebsitesURLLimit, ctl.Probe.Config().Nettests.WebsitesEnabledCategoryCodes)
if err != nil {
return err
}
ctl.SetInputIdxMap(urlIDMap)
builder, err := ctl.Session.NewExperimentBuilder(
"web_connectivity",
)
if err != nil {
return err
}
return ctl.Run(builder, urls)
}