diff --git a/end-to-end/organization_token_test.go b/end-to-end/organization_token_test.go index dbb7e9cc2..235f37c2c 100644 --- a/end-to-end/organization_token_test.go +++ b/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 +} diff --git a/mocks/landscape/landscapemockservice/landscape_mock_service.go b/mocks/landscape/landscapemockservice/landscape_mock_service.go index b02bd35e7..5da9091c7 100644 --- a/mocks/landscape/landscapemockservice/landscape_mock_service.go +++ b/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)