feat(session): expose CheckIn method (#266)

* feat(session): expose CheckIn method

It seems to me the right thing to do is to query the CheckIn API
from the Session rather than querying it from InputLoader.

Then, InputLoader could just take a reference to a Session-like
interface that allows this functionality.

So, this diff exposes the Session.CheckIn method.

Doing that, in turn, required some refactoring to allow for
more and better unit tests.

While doing that, I also noticed that Session required a mutex
to be a well-behaving type, so I did that.

While doing that, I also tried to cover all the lines in session.go
and, as part of that, I have removed unused code.

Reference issue: https://github.com/ooni/probe/issues/1299.

* fix: reinstate comment I shan't have removed

* fix: repair broken test

* fix: a bit more coverage, annotations, etc.

* Update internal/engine/session.go

* Update internal/engine/session_integration_test.go

* Update internal/engine/session_internal_test.go

internal/engine/experiment.go

@@ -1,4 +1,4 @@
-// Package engine contains the engine API
+// Package engine contains the engine API.
 package engine
 
 import (
@@ -188,10 +188,7 @@ func (e *Experiment) OpenReportContext(ctx context.Context) error {
 			Counter:      e.byteCounter,
 		},
 	}
-	if e.session.selectedProbeService == nil {
-		return errors.New("no probe services selected")
-	}
-	client, err := probeservices.NewClient(e.session, *e.session.selectedProbeService)
+	client, err := e.session.NewProbeServicesClient(ctx)
 	if err != nil {
 		e.session.logger.Debugf("%+v", err)
 		return err

internal/engine/session.go

@@ -40,7 +40,7 @@ type SessionConfig struct {
 	TorBinary              string
 }
 
-// Session is a measurement session
+// Session is a measurement session.
 type Session struct {
 	assetsDir                string
 	availableProbeServices   []model.Service
@@ -63,6 +63,32 @@ type Session struct {
 	tunnelMu                 sync.Mutex
 	tunnelName               string
 	tunnel                   tunnel.Tunnel
+
+	// mu provides mutual exclusion.
+	mu sync.Mutex
+
+	// testLookupLocationContext is a an optional hook for testing
+	// allowing us to mock LookupLocationContext.
+	testLookupLocationContext func(ctx context.Context) (*geolocate.Results, error)
+
+	// testMaybeLookupBackendsContext is an optional hook for testing
+	// allowing us to mock MaybeLookupBackendsContext.
+	testMaybeLookupBackendsContext func(ctx context.Context) error
+
+	// testMaybeLookupLocationContext is an optional hook for testing
+	// allowing us to mock MaybeLookupLocationContext.
+	testMaybeLookupLocationContext func(ctx context.Context) error
+
+	// testNewProbeServicesClientForCheckIn is an optional hook for testing
+	// allowing us to mock NewProbeServicesClient when calling CheckIn.
+	testNewProbeServicesClientForCheckIn func(ctx context.Context) (
+		sessionProbeServicesClientForCheckIn, error)
+}
+
+// sessionProbeServicesClientForCheckIn returns the probe services
+// client that we should be using for performing the check-in.
+type sessionProbeServicesClientForCheckIn interface {
+	CheckIn(ctx context.Context, config model.CheckInConfig) (*model.CheckInInfo, error)
 }
 
 // NewSession creates a new session or returns an error
@@ -138,12 +164,98 @@ func (s *Session) KibiBytesSent() float64 {
 	return s.byteCounter.KibiBytesSent()
 }
 
+// CheckIn calls the check-in API. The input arguments MUST NOT
+// be nil. Before querying the API, this function will ensure
+// that the config structure does not contain any field that
+// SHOULD be initialized and is not initialized. Whenever there
+// is a field that is not initialized, we will attempt to set
+// a reasonable default value for such a field. This list describes
+// the current defaults we'll choose:
+//
+// - Platform: if empty, set to Session.Platform();
+//
+// - ProbeASN: if empty, set to Session.ProbeASNString();
+//
+// - ProbeCC: if empty, set to Session.ProbeCC();
+//
+// - RunType: if empty, set to "timed";
+//
+// - SoftwareName: if empty, set to Session.SoftwareName();
+//
+// - SoftwareVersion: if empty, set to Session.SoftwareVersion();
+//
+// - WebConnectivity.CategoryCodes: if nil, we will allocate
+// an empty array (the API does not like nil).
+//
+// Because we MAY need to know the current ASN and CC, this
+// function MAY call MaybeLookupLocationContext.
+//
+// The return value is either the check-in response or an error.
+func (s *Session) CheckIn(
+	ctx context.Context, config *model.CheckInConfig) (*model.CheckInInfo, error) {
+	if err := s.maybeLookupLocationContext(ctx); err != nil {
+		return nil, err
+	}
+	client, err := s.newProbeServicesClientForCheckIn(ctx)
+	if err != nil {
+		return nil, err
+	}
+	if config.Platform == "" {
+		config.Platform = s.Platform()
+	}
+	if config.ProbeASN == "" {
+		config.ProbeASN = s.ProbeASNString()
+	}
+	if config.ProbeCC == "" {
+		config.ProbeCC = s.ProbeCC()
+	}
+	if config.RunType == "" {
+		config.RunType = "timed" // most conservative choice
+	}
+	if config.SoftwareName == "" {
+		config.SoftwareName = s.SoftwareName()
+	}
+	if config.SoftwareVersion == "" {
+		config.SoftwareVersion = s.SoftwareVersion()
+	}
+	if config.WebConnectivity.CategoryCodes == nil {
+		config.WebConnectivity.CategoryCodes = []string{}
+	}
+	return client.CheckIn(ctx, *config)
+}
+
+// maybeLookupLocationContext is a wrapper for MaybeLookupLocationContext that calls
+// the configurable testMaybeLookupLocationContext mock, if configured, and the
+// real MaybeLookupLocationContext API otherwise.
+func (s *Session) maybeLookupLocationContext(ctx context.Context) error {
+	if s.testMaybeLookupLocationContext != nil {
+		return s.testMaybeLookupLocationContext(ctx)
+	}
+	return s.MaybeLookupLocationContext(ctx)
+}
+
+// newProbeServicesClientForCheckIn is a wrapper for NewProbeServicesClientForCheckIn
+// that calls the configurable testNewProbeServicesClientForCheckIn mock, if
+// configured, and the real NewProbeServicesClient API otherwise.
+func (s *Session) newProbeServicesClientForCheckIn(
+	ctx context.Context) (sessionProbeServicesClientForCheckIn, error) {
+	if s.testNewProbeServicesClientForCheckIn != nil {
+		return s.testNewProbeServicesClientForCheckIn(ctx)
+	}
+	client, err := s.NewProbeServicesClient(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return client, nil
+}
+
 // Close ensures that we close all the idle connections that the HTTP clients
 // we are currently using may have created. It will also remove the temp dir
 // that contains data from this session. Not calling this function may likely
 // cause memory leaks in your application because of open idle connections,
 // as well as excessive usage of disk space.
 func (s *Session) Close() error {
+	// TODO(bassosimone): introduce a sync.Once to make this method idempotent.
 	s.httpDefaultTransport.CloseIdleConnections()
 	s.resolver.CloseIdleConnections()
 	s.logger.Infof("%s", s.resolver.Stats())
@@ -161,6 +273,8 @@ func (s *Session) CountryDatabasePath() string {
 // GetTestHelpersByName returns the available test helpers that
 // use the specified name, or false if there's none.
 func (s *Session) GetTestHelpersByName(name string) ([]model.Service, bool) {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	services, ok := s.availableTestHelpers[name]
 	return services, ok
 }
@@ -187,12 +301,7 @@ func (s *Session) MaybeLookupLocation() error {
 
 // MaybeLookupBackends is a caching OONI backends lookup call.
 func (s *Session) MaybeLookupBackends() error {
-	return s.maybeLookupBackends(context.Background())
-}
-
-// MaybeLookupBackendsContext is like MaybeLookupBackends but with context.
-func (s *Session) MaybeLookupBackendsContext(ctx context.Context) (err error) {
-	return s.maybeLookupBackends(ctx)
+	return s.MaybeLookupBackendsContext(context.Background())
 }
 
 // ErrAlreadyUsingProxy indicates that we cannot create a tunnel with
@@ -213,6 +322,7 @@ var ErrAlreadyUsingProxy = errors.New(
 //
 // The tunnel will be closed by session.Close().
 func (s *Session) MaybeStartTunnel(ctx context.Context, name string) error {
+	// TODO(bassosimone): see if we can unify tunnelMu and mu.
 	s.tunnelMu.Lock()
 	defer s.tunnelMu.Unlock()
 	if s.tunnel != nil && s.tunnelName == name {
@@ -258,11 +368,15 @@ func (s *Session) NewExperimentBuilder(name string) (*ExperimentBuilder, error)
 // OONI probe services. This function will benchmark the available
 // probe services, and select the fastest. In case all probe services
 // seem to be down, we try again applying circumvention tactics.
+// This function will fail IMMEDIATELY if given a cancelled context.
 func (s *Session) NewProbeServicesClient(ctx context.Context) (*probeservices.Client, error) {
-	if err := s.maybeLookupBackends(ctx); err != nil {
+	if ctx.Err() != nil {
+		return nil, ctx.Err() // helps with testing
+	}
+	if err := s.maybeLookupBackendsContext(ctx); err != nil {
 		return nil, err
 	}
-	if err := s.MaybeLookupLocationContext(ctx); err != nil {
+	if err := s.maybeLookupLocationContext(ctx); err != nil {
 		return nil, err
 	}
 	if s.selectedProbeServiceHook != nil {
@@ -313,6 +427,8 @@ func (s *Session) ProbeASNString() string {
 
 // ProbeASN returns the probe ASN as an integer.
 func (s *Session) ProbeASN() uint {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	asn := geolocate.DefaultProbeASN
 	if s.location != nil {
 		asn = s.location.ASN
@@ -322,6 +438,8 @@ func (s *Session) ProbeASN() uint {
 
 // ProbeCC returns the probe CC.
 func (s *Session) ProbeCC() string {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	cc := geolocate.DefaultProbeCC
 	if s.location != nil {
 		cc = s.location.CountryCode
@@ -331,6 +449,8 @@ func (s *Session) ProbeCC() string {
 
 // ProbeNetworkName returns the probe network name.
 func (s *Session) ProbeNetworkName() string {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	nn := geolocate.DefaultProbeNetworkName
 	if s.location != nil {
 		nn = s.location.NetworkName
@@ -340,6 +460,8 @@ func (s *Session) ProbeNetworkName() string {
 
 // ProbeIP returns the probe IP.
 func (s *Session) ProbeIP() string {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	ip := geolocate.DefaultProbeIP
 	if s.location != nil {
 		ip = s.location.ProbeIP
@@ -359,6 +481,8 @@ func (s *Session) ResolverASNString() string {
 
 // ResolverASN returns the resolver ASN
 func (s *Session) ResolverASN() uint {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	asn := geolocate.DefaultResolverASN
 	if s.location != nil {
 		asn = s.location.ResolverASN
@@ -368,6 +492,8 @@ func (s *Session) ResolverASN() uint {
 
 // ResolverIP returns the resolver IP
 func (s *Session) ResolverIP() string {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	ip := geolocate.DefaultResolverIP
 	if s.location != nil {
 		ip = s.location.ResolverIP
@@ -377,6 +503,8 @@ func (s *Session) ResolverIP() string {
 
 // ResolverNetworkName returns the resolver network name.
 func (s *Session) ResolverNetworkName() string {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	nn := geolocate.DefaultResolverNetworkName
 	if s.location != nil {
 		nn = s.location.ResolverNetworkName
@@ -423,7 +551,10 @@ func (s *Session) MaybeUpdateResources(ctx context.Context) error {
 	return (&resourcesmanager.CopyWorker{DestDir: s.assetsDir}).Ensure()
 }
 
-func (s *Session) getAvailableProbeServices() []model.Service {
+// getAvailableProbeServicesUnlocked returns the available probe
+// services. This function WILL NOT acquire the mu mutex, therefore,
+// you MUST ensure you are using it from a locked context.
+func (s *Session) getAvailableProbeServicesUnlocked() []model.Service {
 	if len(s.availableProbeServices) > 0 {
 		return s.availableProbeServices
 	}
@@ -458,22 +589,27 @@ func (s *Session) initOrchestraClient(
 	return clnt, nil
 }
 
-// LookupASN maps an IP address to its ASN and network name. This method implements
-// LocationLookupASNLookupper.LookupASN.
-func (s *Session) LookupASN(dbPath, ip string) (uint, string, error) {
-	return geolocate.LookupASN(dbPath, ip)
-}
-
 // ErrAllProbeServicesFailed indicates all probe services failed.
 var ErrAllProbeServicesFailed = errors.New("all available probe services failed")
 
-func (s *Session) maybeLookupBackends(ctx context.Context) error {
-	// TODO(bassosimone): do we need a mutex here?
+// maybeLookupBackendsContext uses testMaybeLookupBackendsContext if
+// not nil, otherwise it calls MaybeLookupBackendsContext.
+func (s *Session) maybeLookupBackendsContext(ctx context.Context) error {
+	if s.testMaybeLookupBackendsContext != nil {
+		return s.testMaybeLookupBackendsContext(ctx)
+	}
+	return s.MaybeLookupBackendsContext(ctx)
+}
+
+// MaybeLookupBackendsContext is like MaybeLookupBackends but with context.
+func (s *Session) MaybeLookupBackendsContext(ctx context.Context) error {
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	if s.selectedProbeService != nil {
 		return nil
 	}
 	s.queryProbeServicesCount.Add(1)
-	candidates := probeservices.TryAll(ctx, s, s.getAvailableProbeServices())
+	candidates := probeservices.TryAll(ctx, s, s.getAvailableProbeServicesUnlocked())
 	selected := probeservices.SelectBest(candidates)
 	if selected == nil {
 		return ErrAllProbeServicesFailed
@@ -499,11 +635,26 @@ func (s *Session) LookupLocationContext(ctx context.Context) (*geolocate.Results
 	return task.Run(ctx)
 }
 
+// lookupLocationContext calls testLookupLocationContext if set and
+// otherwise calls LookupLocationContext.
+func (s *Session) lookupLocationContext(ctx context.Context) (*geolocate.Results, error) {
+	if s.testLookupLocationContext != nil {
+		return s.testLookupLocationContext(ctx)
+	}
+	return s.LookupLocationContext(ctx)
+}
+
 // MaybeLookupLocationContext is like MaybeLookupLocation but with a context
-// that can be used to interrupt this long running operation.
+// that can be used to interrupt this long running operation. This function
+// will fail IMMEDIATELY if given a cancelled context.
 func (s *Session) MaybeLookupLocationContext(ctx context.Context) error {
+	if ctx.Err() != nil {
+		return ctx.Err() // helps with testing
+	}
+	defer s.mu.Unlock()
+	s.mu.Lock()
 	if s.location == nil {
-		location, err := s.LookupLocationContext(ctx)
+		location, err := s.lookupLocationContext(ctx)
 		if err != nil {
 			return err
 		}

internal/engine/session_integration_test.go

@@ -3,6 +3,7 @@ package engine
 import (
 	"context"
 	"errors"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -10,17 +11,34 @@ import (
 	"os"
 	"syscall"
 	"testing"
-	"time"
 
 	"github.com/apex/log"
 	"github.com/google/go-cmp/cmp"
 	"github.com/ooni/probe-cli/v3/internal/engine/geolocate"
 	"github.com/ooni/probe-cli/v3/internal/engine/model"
-	"github.com/ooni/probe-cli/v3/internal/engine/netx"
 	"github.com/ooni/probe-cli/v3/internal/engine/probeservices"
 	"github.com/ooni/probe-cli/v3/internal/version"
 )
 
+func TestSessionByteCounter(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skip test in short mode")
+	}
+	s := newSessionForTesting(t)
+	client := s.DefaultHTTPClient()
+	resp, err := client.Get("https://www.google.com")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer resp.Body.Close()
+	if _, err := io.Copy(ioutil.Discard, resp.Body); err != nil {
+		t.Fatal(err)
+	}
+	if s.KibiBytesSent() <= 0 || s.KibiBytesReceived() <= 0 {
+		t.Fatal("byte counter is not working")
+	}
+}
+
 func TestNewSessionBuilderChecks(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skip test in short mode")
@@ -307,6 +325,21 @@ func TestSessionLocationLookup(t *testing.T) {
 	}
 }
 
+func TestSessionCheckInWithRealAPI(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skip test in short mode")
+	}
+	sess := newSessionForTesting(t)
+	defer sess.Close()
+	results, err := sess.CheckIn(context.Background(), &model.CheckInConfig{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if results == nil {
+		t.Fatal("expected non nil results here")
+	}
+}
+
 func TestSessionCloseCancelsTempDir(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skip test in short mode")
@@ -584,50 +617,31 @@ func TestNewOrchestraClientMaybeLookupBackendsFailure(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skip test in short mode")
 	}
+	errMocked := errors.New("mocked error")
 	sess := newSessionForTestingNoLookups(t)
-	ctx, cancel := context.WithCancel(context.Background())
-	cancel() // fail immediately
-	client, err := sess.NewOrchestraClient(ctx)
-	if !errors.Is(err, ErrAllProbeServicesFailed) {
-		t.Fatal("not the error we expected")
+	sess.testMaybeLookupBackendsContext = func(ctx context.Context) error {
+		return errMocked
+	}
+	client, err := sess.NewOrchestraClient(context.Background())
+	if !errors.Is(err, errMocked) {
+		t.Fatal("not the error we expected", err)
 	}
 	if client != nil {
 		t.Fatal("expected nil client here")
 	}
 }
 
-type httpTransportThatSleeps struct {
-	txp netx.HTTPRoundTripper
-	st  time.Duration
-}
-
-func (txp httpTransportThatSleeps) RoundTrip(req *http.Request) (*http.Response, error) {
-	resp, err := txp.txp.RoundTrip(req)
-	time.Sleep(txp.st)
-	return resp, err
-}
-
-func (txp httpTransportThatSleeps) CloseIdleConnections() {
-	txp.txp.CloseIdleConnections()
-}
-
 func TestNewOrchestraClientMaybeLookupLocationFailure(t *testing.T) {
 	if testing.Short() {
 		t.Skip("skip test in short mode")
 	}
+	errMocked := errors.New("mocked error")
 	sess := newSessionForTestingNoLookups(t)
-	sess.httpDefaultTransport = httpTransportThatSleeps{
-		txp: sess.httpDefaultTransport,
-		st:  5 * time.Second,
+	sess.testMaybeLookupLocationContext = func(ctx context.Context) error {
+		return errMocked
 	}
-	// The transport sleeps for five seconds, so the context should be expired by
-	// the time in which we attempt at looking up the location. Because the
-	// implementation performs the round-trip and _then_ sleeps, it means we'll
-	// see the context expired error when performing the location lookup.
-	ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
-	defer cancel()
-	client, err := sess.NewOrchestraClient(ctx)
-	if !errors.Is(err, geolocate.ErrAllIPLookuppersFailed) {
+	client, err := sess.NewOrchestraClient(context.Background())
+	if !errors.Is(err, errMocked) {
 		t.Fatalf("not the error we expected: %+v", err)
 	}
 	if client != nil {
@@ -651,3 +665,14 @@ func TestNewOrchestraClientProbeServicesNewClientFailure(t *testing.T) {
 		t.Fatal("expected nil client here")
 	}
 }
+
+func TestSessionNewSubmitterReturnsNonNilSubmitter(t *testing.T) {
+	sess := newSessionForTesting(t)
+	subm, err := sess.NewSubmitter(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if subm == nil {
+		t.Fatal("expected non nil submitter here")
+	}
+}

internal/engine/session_internal_test.go

@@ -1,6 +1,13 @@
 package engine
 
 import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/ooni/probe-cli/v3/internal/engine/geolocate"
 	"github.com/ooni/probe-cli/v3/internal/engine/model"
 )
 
@@ -9,7 +16,7 @@ func (s *Session) SetAssetsDir(assetsDir string) {
 }
 
 func (s *Session) GetAvailableProbeServices() []model.Service {
-	return s.getAvailableProbeServices()
+	return s.getAvailableProbeServicesUnlocked()
 }
 
 func (s *Session) AppendAvailableProbeService(svc model.Service) {
@@ -19,3 +26,189 @@ func (s *Session) AppendAvailableProbeService(svc model.Service) {
 func (s *Session) QueryProbeServicesCount() int64 {
 	return s.queryProbeServicesCount.Load()
 }
+
+// mockableProbeServicesClientForCheckIn allows us to mock the
+// probeservices.Client used by Session.CheckIn.
+type mockableProbeServicesClientForCheckIn struct {
+	// Config is the config passed to the call.
+	Config *model.CheckInConfig
+
+	// Results contains the results of the call. This field MUST be
+	// non-nil if and only if Error is nil.
+	Results *model.CheckInInfo
+
+	// Error indicates whether the call failed. This field MUST be
+	// non-nil if and only if Error is nil.
+	Error error
+
+	// mu provides mutual exclusion.
+	mu sync.Mutex
+}
+
+// CheckIn implements sessionProbeServicesClientForCheckIn.CheckIn.
+func (c *mockableProbeServicesClientForCheckIn) CheckIn(
+	ctx context.Context, config model.CheckInConfig) (*model.CheckInInfo, error) {
+	defer c.mu.Unlock()
+	c.mu.Lock()
+	if c.Config != nil {
+		return nil, errors.New("called more than once")
+	}
+	c.Config = &config
+	if c.Results == nil && c.Error == nil {
+		return nil, errors.New("misconfigured mockableProbeServicesClientForCheckIn")
+	}
+	return c.Results, c.Error
+}
+
+func TestSessionCheckInSuccessful(t *testing.T) {
+	results := &model.CheckInInfo{
+		WebConnectivity: &model.CheckInInfoWebConnectivity{
+			ReportID: "xxx-x-xx",
+			URLs: []model.URLInfo{{
+				CategoryCode: "NEWS",
+				CountryCode:  "IT",
+				URL:          "https://www.repubblica.it/",
+			}, {
+				CategoryCode: "NEWS",
+				CountryCode:  "IT",
+				URL:          "https://www.unita.it/",
+			}},
+		},
+	}
+	mockedClnt := &mockableProbeServicesClientForCheckIn{
+		Results: results,
+	}
+	s := &Session{
+		location: &geolocate.Results{
+			ASN:         137,
+			CountryCode: "IT",
+		},
+		softwareName:    "miniooni",
+		softwareVersion: "0.1.0-dev",
+		testMaybeLookupLocationContext: func(ctx context.Context) error {
+			return nil
+		},
+		testNewProbeServicesClientForCheckIn: func(
+			ctx context.Context) (sessionProbeServicesClientForCheckIn, error) {
+			return mockedClnt, nil
+		},
+	}
+	out, err := s.CheckIn(context.Background(), &model.CheckInConfig{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if diff := cmp.Diff(results, out); diff != "" {
+		t.Fatal(diff)
+	}
+	if mockedClnt.Config.Platform != s.Platform() {
+		t.Fatal("invalid Config.Platform")
+	}
+	if mockedClnt.Config.ProbeASN != "AS137" {
+		t.Fatal("invalid Config.ProbeASN")
+	}
+	if mockedClnt.Config.ProbeCC != "IT" {
+		t.Fatal("invalid Config.ProbeCC")
+	}
+	if mockedClnt.Config.RunType != "timed" {
+		t.Fatal("invalid Config.RunType")
+	}
+	if mockedClnt.Config.SoftwareName != "miniooni" {
+		t.Fatal("invalid Config.SoftwareName")
+	}
+	if mockedClnt.Config.SoftwareVersion != "0.1.0-dev" {
+		t.Fatal("invalid Config.SoftwareVersion")
+	}
+	if mockedClnt.Config.WebConnectivity.CategoryCodes == nil {
+		t.Fatal("invalid ...CategoryCodes")
+	}
+}
+
+func TestSessionCheckInCannotLookupLocation(t *testing.T) {
+	errMocked := errors.New("mocked error")
+	s := &Session{
+		testMaybeLookupLocationContext: func(ctx context.Context) error {
+			return errMocked
+		},
+	}
+	out, err := s.CheckIn(context.Background(), &model.CheckInConfig{})
+	if !errors.Is(err, errMocked) {
+		t.Fatal("no the error we expected", err)
+	}
+	if out != nil {
+		t.Fatal("expected nil result here")
+	}
+}
+
+func TestSessionCheckInCannotCreateProbeServicesClient(t *testing.T) {
+	errMocked := errors.New("mocked error")
+	s := &Session{
+		location: &geolocate.Results{
+			ASN:         137,
+			CountryCode: "IT",
+		},
+		softwareName:    "miniooni",
+		softwareVersion: "0.1.0-dev",
+		testMaybeLookupLocationContext: func(ctx context.Context) error {
+			return nil
+		},
+		testNewProbeServicesClientForCheckIn: func(
+			ctx context.Context) (sessionProbeServicesClientForCheckIn, error) {
+			return nil, errMocked
+		},
+	}
+	out, err := s.CheckIn(context.Background(), &model.CheckInConfig{})
+	if !errors.Is(err, errMocked) {
+		t.Fatal("no the error we expected", err)
+	}
+	if out != nil {
+		t.Fatal("expected nil result here")
+	}
+}
+
+func TestLowercaseMaybeLookupLocationContextWithCancelledContext(t *testing.T) {
+	s := &Session{}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // immediately kill the context
+	err := s.maybeLookupLocationContext(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatal("not the error we expected", err)
+	}
+}
+
+func TestNewProbeServicesClientForCheckIn(t *testing.T) {
+	s := &Session{}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // immediately kill the context
+	clnt, err := s.newProbeServicesClientForCheckIn(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatal("not the error we expected", err)
+	}
+	if clnt != nil {
+		t.Fatal("expected nil client here")
+	}
+}
+
+func TestSessionNewSubmitterWithCancelledContext(t *testing.T) {
+	sess := newSessionForTesting(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // fail immediately
+	subm, err := sess.NewSubmitter(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatal("not the error we expected", err)
+	}
+	if subm != nil {
+		t.Fatal("expected nil submitter here")
+	}
+}
+
+func TestSessionMaybeLookupLocationContextLookupLocationContextFailure(t *testing.T) {
+	errMocked := errors.New("mocked error")
+	sess := newSessionForTestingNoLookups(t)
+	sess.testLookupLocationContext = func(ctx context.Context) (*geolocate.Results, error) {
+		return nil, errMocked
+	}
+	err := sess.MaybeLookupLocationContext(context.Background())
+	if !errors.Is(err, errMocked) {
+		t.Fatal("not the error we expected", err)
+	}
+}

pkg/oonimkall/session.go

@@ -479,17 +479,17 @@ func (sess *Session) FetchURLList(ctx *Context, config *URLListConfig) (*URLList
 	if config.CountryCode == "" {
 		config.CountryCode = "XX"
 		info, err := sess.sessp.LookupLocationContext(ctx.ctx)
+		// TODO(bassosimone): this piece of code feels wrong to me. We don't
+		// want to continue if we cannot discover the country.
 		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

pkg/oonimkall/session_integration_test.go

@@ -12,7 +12,6 @@ import (
 	"testing"
 	"time"
 
-	engine "github.com/ooni/probe-cli/v3/internal/engine"
 	"github.com/ooni/probe-cli/v3/internal/engine/geolocate"
 	"github.com/ooni/probe-cli/v3/internal/engine/model"
 	"github.com/ooni/probe-cli/v3/pkg/oonimkall"
@@ -362,7 +361,7 @@ func TestCheckInNewProbeServicesFailure(t *testing.T) {
 	config.WebConnectivity.Add("NEWS")
 	config.WebConnectivity.Add("CULTR")
 	result, err := sess.CheckIn(ctx, &config)
-	if !errors.Is(err, engine.ErrAllProbeServicesFailed) {
+	if !errors.Is(err, context.Canceled) {
 		t.Fatalf("not the error we expected: %+v", err)
 	}
 	if result != nil {
@@ -440,11 +439,18 @@ func TestFetchURLListSuccess(t *testing.T) {
 	if result == nil || result.Results == nil {
 		t.Fatal("got nil result")
 	}
-	for _, entry := range result.Results {
+	for idx := int64(0); idx < result.Size(); idx++ {
+		entry := result.At(idx)
 		if entry.CategoryCode != "NEWS" && entry.CategoryCode != "CULTR" {
 			t.Fatalf("unexpected category code: %+v", entry)
 		}
 	}
+	if result.At(-1) != nil {
+		t.Fatal("expected nil here")
+	}
+	if result.At(result.Size()) != nil {
+		t.Fatal("expected nil here")
+	}
 }
 
 func TestFetchURLListWithCC(t *testing.T) {

pkg/oonimkall/task_integration_test.go

@@ -58,10 +58,10 @@ func TestGood(t *testing.T) {
 		if err := json.Unmarshal([]byte(eventstr), &event); err != nil {
 			t.Fatal(err)
 		}
-		if event.Key != "task_terminated" {
-			t.Fatalf("unexpected event.Key: %s", event.Key)
+		if event.Key == "task_terminated" {
+			break
 		}
-		break
+		t.Fatalf("unexpected event.Key: %s", event.Key)
 	}
 }