Added landscape checks to TestOrganizationProvidedToken

end-to-end/organization_token_test.go

@@ -4,13 +4,17 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"net"
 	"os"
 	"testing"
 	"time"
 
+	landscapeapi "github.com/canonical/landscape-hostagent-api"
+	"github.com/canonical/ubuntu-pro-for-windows/mocks/landscape/landscapemockservice"
 	"github.com/stretchr/testify/require"
 	wsl "github.com/ubuntu/gowsl"
 	"golang.org/x/sys/windows/registry"
+	"google.golang.org/grpc"
 )
 
 func TestOrganizationProvidedToken(t *testing.T) {
@@ -35,12 +39,22 @@ func TestOrganizationProvidedToken(t *testing.T) {
 	for name, tc := range testCases {
 		tc := tc
 		t.Run(name, func(t *testing.T) {
-			ctx := context.Background()
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
 
 			testSetup(t)
 
+			lis, server, landscapeService := setUpLandscapeMock(t, ctx, "localhost:")
+			setUbuntuProRegistry(t, "LandscapeAgentURL", lis.Addr().String())
+
+			go server.Serve(lis)
+			defer server.Stop()
+
+			proToken := os.Getenv(proTokenEnv)
+			require.NotEmptyf(t, proToken, "Setup: environment variable %q should contain a valid pro token, but is empty", proTokenEnv)
+			setUbuntuProRegistry(t, "ProTokenOrg", proToken)
+
 			if tc.whenToken == beforeDistroRegistration {
-				activateOrgSubscription(t)
 				cleanup := startAgent(t, ctx)
 				defer cleanup()
 			}
@@ -58,7 +72,6 @@ func TestOrganizationProvidedToken(t *testing.T) {
 				err := d.Terminate()
 				require.NoError(t, err, "could not restart distro")
 
-				activateOrgSubscription(t)
 				cleanup := startAgent(t, ctx)
 				defer cleanup()
 
@@ -83,6 +96,13 @@ func TestOrganizationProvidedToken(t *testing.T) {
 				}
 				return attached
 			}, maxTimeout, time.Second, "distro should have been Pro attached")
+
+			for _, info := range landscapeService.Hosts() {
+				if info.Token == proToken {
+					return
+				}
+			}
+			require.Fail(t, "Landscape connection did not happen")
 		})
 	}
 }
@@ -108,17 +128,14 @@ func distroIsProAttached(t *testing.T, d wsl.Distro) (bool, error) {
 	return response.Attached, nil
 }
 
-func activateOrgSubscription(t *testing.T) {
+func setUbuntuProRegistry(t *testing.T, field string, value string) {
 	t.Helper()
 
-	token := os.Getenv(proTokenEnv)
-	require.NotEmptyf(t, token, "Setup: environment variable %q should contain a valid pro token, but is empty", proTokenEnv)
-
 	key, _, err := registry.CreateKey(registry.CURRENT_USER, registryPath, registry.WRITE)
 	require.NoErrorf(t, err, "Setup: could not open UbuntuPro registry key")
 	defer key.Close()
 
-	err = key.SetStringValue("ProTokenOrg", token)
+	err = key.SetStringValue(field, value)
 	require.NoError(t, err, "could not write token in registry")
 }
 
@@ -132,3 +149,18 @@ func logWslProServiceJournal(t *testing.T, ctx context.Context, d wsl.Distro) {
 	}
 	t.Logf("wsl-pro-service logs:\n%s\n", out)
 }
+
+//nolint:revive // Context goes after testing.T
+func setUpLandscapeMock(t *testing.T, ctx context.Context, addr string) (lis net.Listener, server *grpc.Server, service *landscapemockservice.Service) {
+	t.Helper()
+
+	var cfg net.ListenConfig
+	lis, err := cfg.Listen(ctx, "tcp", addr)
+	require.NoError(t, err, "Setup: can't listen")
+
+	server = grpc.NewServer()
+	service = landscapemockservice.New()
+	landscapeapi.RegisterLandscapeHostAgentServer(server, service)
+
+	return lis, server, service
+}

mocks/landscape/landscapemockservice/landscape_mock_service.go

@@ -111,6 +111,11 @@ func (s *Service) Connect(stream landscapeapi.LandscapeHostAgent_ConnectServer)
 			return nil
 		}
 
+		if hostInfo.Token == "" {
+			// No token: no Landscape
+			continue
+		}
+
 		s.mu.Lock()
 
 		s.recvLog = append(s.recvLog, hostInfo)