github.com/gravitational/teleport/api@v0.0.0-20240507183017-3110591cbafc/client/client_test.go (about)

     1  /*
     2  Copyright 2021 Gravitational, Inc.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package client
    18  
    19  import (
    20  	"context"
    21  	"flag"
    22  	"fmt"
    23  	"net"
    24  	"os"
    25  	"strings"
    26  	"sync/atomic"
    27  	"testing"
    28  	"time"
    29  
    30  	"github.com/google/go-cmp/cmp"
    31  	"github.com/gravitational/trace"
    32  	"github.com/gravitational/trace/trail"
    33  	"github.com/stretchr/testify/assert"
    34  	"github.com/stretchr/testify/require"
    35  
    36  	"github.com/gravitational/teleport/api"
    37  	"github.com/gravitational/teleport/api/client/proto"
    38  	"github.com/gravitational/teleport/api/defaults"
    39  	"github.com/gravitational/teleport/api/metadata"
    40  	"github.com/gravitational/teleport/api/types"
    41  )
    42  
    43  func TestMain(m *testing.M) {
    44  	flag.Parse()
    45  	os.Exit(m.Run())
    46  }
    47  
    48  type pingService struct {
    49  	*proto.UnimplementedAuthServiceServer
    50  	userAgentFromLastCallValue atomic.Value
    51  }
    52  
    53  func (s *pingService) Ping(ctx context.Context, req *proto.PingRequest) (*proto.PingResponse, error) {
    54  	s.userAgentFromLastCallValue.Store(metadata.UserAgentFromContext(ctx))
    55  	return &proto.PingResponse{}, nil
    56  }
    57  
    58  func (s *pingService) userAgentFromLastCall() string {
    59  	if userAgent, ok := s.userAgentFromLastCallValue.Load().(string); ok {
    60  		return userAgent
    61  	}
    62  	return ""
    63  }
    64  
    65  func TestNew(t *testing.T) {
    66  	t.Parallel()
    67  	ctx := context.Background()
    68  	srv := startMockServer(t, &pingService{})
    69  
    70  	tests := []struct {
    71  		desc         string
    72  		modifyConfig func(*Config)
    73  		assertErr    require.ErrorAssertionFunc
    74  	}{{
    75  		desc:         "successfully dial tcp address.",
    76  		modifyConfig: func(c *Config) { /* noop */ },
    77  		assertErr:    require.NoError,
    78  	}, {
    79  		desc: "synchronously dial addr/cred pairs and succeed with the 1 good pair.",
    80  		modifyConfig: func(c *Config) {
    81  			c.Addrs = append(c.Addrs, "bad addr", "bad addr")
    82  			c.Credentials = append([]Credentials{&tlsConfigCreds{nil}, &tlsConfigCreds{nil}}, c.Credentials...)
    83  		},
    84  		assertErr: require.NoError,
    85  	}, {
    86  		desc: "fail to dial with a bad address.",
    87  		modifyConfig: func(c *Config) {
    88  			c.Addrs = []string{"bad addr"}
    89  		},
    90  		assertErr: func(t require.TestingT, err error, _ ...interface{}) {
    91  			require.Error(t, err)
    92  			require.ErrorContains(t, err, "all connection methods failed")
    93  		},
    94  	}, {
    95  		desc: "fail to dial with no address or dialer.",
    96  		modifyConfig: func(c *Config) {
    97  			c.Addrs = nil
    98  		},
    99  		assertErr: func(t require.TestingT, err error, _ ...interface{}) {
   100  			require.Error(t, err)
   101  			require.ErrorContains(t, err, "no connection methods found, try providing Dialer or Addrs in config")
   102  		},
   103  	}}
   104  
   105  	for _, tt := range tests {
   106  		t.Run(tt.desc, func(t *testing.T) {
   107  			cfg := srv.clientCfg()
   108  			tt.modifyConfig(&cfg)
   109  
   110  			clt, err := New(ctx, cfg)
   111  			tt.assertErr(t, err)
   112  			if err != nil {
   113  				return
   114  			}
   115  
   116  			// Requests to the server should succeed.
   117  			_, err = clt.Ping(ctx)
   118  			assert.NoError(t, err, "Ping failed")
   119  			assert.NoError(t, clt.Close(), "Close failed")
   120  		})
   121  	}
   122  }
   123  
   124  func TestNewDialBackground(t *testing.T) {
   125  	t.Parallel()
   126  	ctx := context.Background()
   127  
   128  	// Create a server but don't serve it yet.
   129  	l, err := net.Listen("tcp", "localhost:")
   130  	require.NoError(t, err)
   131  	addr := l.Addr().String()
   132  	ping := &pingService{}
   133  	srv := newMockServer(t, addr, ping)
   134  
   135  	// Create client before the server is listening.
   136  	cfg := srv.clientCfg()
   137  	cfg.DialInBackground = true
   138  	cfg.DialOpts = append(cfg.DialOpts, metadata.WithUserAgentFromTeleportComponent("api-client-test"))
   139  	clt, err := New(ctx, cfg)
   140  	require.NoError(t, err)
   141  	t.Cleanup(func() { require.NoError(t, clt.Close()) })
   142  
   143  	// requests to the server will result in a connection error.
   144  	cancelCtx, cancel := context.WithTimeout(ctx, time.Second*3)
   145  	defer cancel()
   146  	_, err = clt.Ping(cancelCtx)
   147  	require.Error(t, err)
   148  
   149  	// Server the listener and wait for the client connection to be ready.
   150  	srv.serve(t, l)
   151  	require.NoError(t, clt.waitForConnectionReady(ctx))
   152  
   153  	// requests to the server should succeed.
   154  	_, err = clt.Ping(ctx)
   155  	require.NoError(t, err)
   156  
   157  	// Verify user agent.
   158  	expectUserAgentPrefix := fmt.Sprintf("api-client-test/%v grpc-go/", api.Version)
   159  	require.True(t, strings.HasPrefix(ping.userAgentFromLastCall(), expectUserAgentPrefix))
   160  }
   161  
   162  func TestWaitForConnectionReady(t *testing.T) {
   163  	t.Parallel()
   164  	ctx := context.Background()
   165  
   166  	// Create a server but don't serve it yet.
   167  	l, err := net.Listen("tcp", "localhost:")
   168  	require.NoError(t, err)
   169  	addr := l.Addr().String()
   170  	srv := newMockServer(t, addr, &proto.UnimplementedAuthServiceServer{})
   171  
   172  	// Create client before the server is listening.
   173  	cfg := srv.clientCfg()
   174  	cfg.DialInBackground = true
   175  	clt, err := New(ctx, cfg)
   176  	require.NoError(t, err)
   177  	t.Cleanup(func() { require.NoError(t, clt.Close()) })
   178  
   179  	// WaitForConnectionReady should return an error once the
   180  	// context is canceled if the server isn't open to connections.
   181  	cancelCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
   182  	defer cancel()
   183  	require.Error(t, clt.waitForConnectionReady(cancelCtx))
   184  
   185  	// WaitForConnectionReady should return nil if the server is open to connections.
   186  	srv.serve(t, l)
   187  	require.NoError(t, clt.waitForConnectionReady(ctx))
   188  
   189  	// WaitForConnectionReady should return an error if the grpc connection is closed.
   190  	require.NoError(t, clt.Close())
   191  	require.Error(t, clt.waitForConnectionReady(ctx))
   192  }
   193  
   194  type listResourcesService struct {
   195  	*proto.UnimplementedAuthServiceServer
   196  }
   197  
   198  func (s *listResourcesService) ListResources(ctx context.Context, req *proto.ListResourcesRequest) (*proto.ListResourcesResponse, error) {
   199  	expectedResources, err := testResources[types.ResourceWithLabels](req.ResourceType, req.Namespace)
   200  	if err != nil {
   201  		return nil, trail.ToGRPC(err)
   202  	}
   203  
   204  	resp := &proto.ListResourcesResponse{
   205  		Resources:  make([]*proto.PaginatedResource, 0, len(expectedResources)),
   206  		TotalCount: int32(len(expectedResources)),
   207  	}
   208  
   209  	var (
   210  		takeResources    = req.StartKey == ""
   211  		lastResourceName string
   212  	)
   213  	for _, resource := range expectedResources {
   214  		if resource.GetName() == req.StartKey {
   215  			takeResources = true
   216  			continue
   217  		}
   218  
   219  		if !takeResources {
   220  			continue
   221  		}
   222  
   223  		var protoResource *proto.PaginatedResource
   224  		switch req.ResourceType {
   225  		case types.KindDatabaseServer:
   226  			database, ok := resource.(*types.DatabaseServerV3)
   227  			if !ok {
   228  				return nil, trace.Errorf("database server has invalid type %T", resource)
   229  			}
   230  
   231  			protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_DatabaseServer{DatabaseServer: database}}
   232  		case types.KindAppServer:
   233  			app, ok := resource.(*types.AppServerV3)
   234  			if !ok {
   235  				return nil, trace.Errorf("application server has invalid type %T", resource)
   236  			}
   237  
   238  			protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_AppServer{AppServer: app}}
   239  		case types.KindNode:
   240  			srv, ok := resource.(*types.ServerV2)
   241  			if !ok {
   242  				return nil, trace.Errorf("node has invalid type %T", resource)
   243  			}
   244  
   245  			protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_Node{Node: srv}}
   246  		case types.KindKubeServer:
   247  			srv, ok := resource.(*types.KubernetesServerV3)
   248  			if !ok {
   249  				return nil, trace.Errorf("kubernetes server has invalid type %T", resource)
   250  			}
   251  
   252  			protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_KubernetesServer{KubernetesServer: srv}}
   253  		case types.KindWindowsDesktop:
   254  			desktop, ok := resource.(*types.WindowsDesktopV3)
   255  			if !ok {
   256  				return nil, trace.Errorf("windows desktop has invalid type %T", resource)
   257  			}
   258  
   259  			protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_WindowsDesktop{WindowsDesktop: desktop}}
   260  		case types.KindAppOrSAMLIdPServiceProvider:
   261  			//nolint:staticcheck // SA1019. TODO(sshah) DELETE IN 17.0
   262  			appServerOrSP, ok := resource.(*types.AppServerOrSAMLIdPServiceProviderV1)
   263  			if !ok {
   264  				return nil, trace.Errorf("AppServerOrSAMLIdPServiceProvider has invalid type %T", resource)
   265  			}
   266  
   267  			protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_AppServerOrSAMLIdPServiceProvider{AppServerOrSAMLIdPServiceProvider: appServerOrSP}}
   268  		case types.KindSAMLIdPServiceProvider:
   269  			samlSP, ok := resource.(*types.SAMLIdPServiceProviderV1)
   270  			if !ok {
   271  				return nil, trace.Errorf("SAML IdP service provider has invalid type %T", resource)
   272  			}
   273  
   274  			protoResource = &proto.PaginatedResource{Resource: &proto.PaginatedResource_SAMLIdPServiceProvider{SAMLIdPServiceProvider: samlSP}}
   275  		}
   276  		resp.Resources = append(resp.Resources, protoResource)
   277  		lastResourceName = resource.GetName()
   278  		if len(resp.Resources) == int(req.Limit) {
   279  			break
   280  		}
   281  	}
   282  
   283  	if len(resp.Resources) != len(expectedResources) {
   284  		resp.NextKey = lastResourceName
   285  	}
   286  
   287  	return resp, nil
   288  }
   289  
   290  const fiveMBNode = "fiveMBNode"
   291  
   292  func testResources[T types.ResourceWithLabels](resourceType, namespace string) ([]T, error) {
   293  	size := 50
   294  	// Artificially make each node ~ 100KB to force
   295  	// ListResources to fail with chunks of >= 40.
   296  	labelSize := 100000
   297  	resources := make([]T, 0, size)
   298  
   299  	switch resourceType {
   300  	case types.KindDatabaseServer:
   301  		for i := 0; i < size; i++ {
   302  			resource, err := types.NewDatabaseServerV3(types.Metadata{
   303  				Name: fmt.Sprintf("db-%d", i),
   304  				Labels: map[string]string{
   305  					"label": string(make([]byte, labelSize)),
   306  				},
   307  			}, types.DatabaseServerSpecV3{
   308  				Hostname: "localhost",
   309  				HostID:   fmt.Sprintf("host-%d", i),
   310  				Database: &types.DatabaseV3{
   311  					Metadata: types.Metadata{
   312  						Name: fmt.Sprintf("db-%d", i),
   313  					},
   314  					Spec: types.DatabaseSpecV3{
   315  						Protocol: types.DatabaseProtocolPostgreSQL,
   316  						URI:      "localhost",
   317  					},
   318  				},
   319  			})
   320  			if err != nil {
   321  				return nil, trace.Wrap(err)
   322  			}
   323  
   324  			resources = append(resources, any(resource).(T))
   325  		}
   326  	case types.KindAppServer:
   327  		for i := 0; i < size; i++ {
   328  			app, err := types.NewAppV3(types.Metadata{
   329  				Name: fmt.Sprintf("app-%d", i),
   330  			}, types.AppSpecV3{
   331  				URI: "localhost",
   332  			})
   333  			if err != nil {
   334  				return nil, trace.Wrap(err)
   335  			}
   336  
   337  			resource, err := types.NewAppServerV3(types.Metadata{
   338  				Name: fmt.Sprintf("app-%d", i),
   339  				Labels: map[string]string{
   340  					"label": string(make([]byte, labelSize)),
   341  				},
   342  			}, types.AppServerSpecV3{
   343  				HostID: fmt.Sprintf("host-%d", i),
   344  				App:    app,
   345  			})
   346  			if err != nil {
   347  				return nil, trace.Wrap(err)
   348  			}
   349  
   350  			resources = append(resources, any(resource).(T))
   351  		}
   352  	case types.KindNode:
   353  		for i := 0; i < size; i++ {
   354  			nodeLabelSize := labelSize
   355  			if namespace == fiveMBNode && i == 0 {
   356  				// Artificially make a node ~ 5MB to force
   357  				// ListNodes to fail regardless of chunk size.
   358  				nodeLabelSize = 5000000
   359  			}
   360  
   361  			var err error
   362  			resource, err := types.NewServerWithLabels(fmt.Sprintf("node-%d", i), types.KindNode, types.ServerSpecV2{},
   363  				map[string]string{
   364  					"label": string(make([]byte, nodeLabelSize)),
   365  				},
   366  			)
   367  			if err != nil {
   368  				return nil, trace.Wrap(err)
   369  			}
   370  
   371  			resources = append(resources, any(resource).(T))
   372  		}
   373  	case types.KindKubeServer:
   374  		for i := 0; i < size; i++ {
   375  			var err error
   376  			name := fmt.Sprintf("kube-service-%d", i)
   377  			kube, err := types.NewKubernetesClusterV3(types.Metadata{
   378  				Name:   name,
   379  				Labels: map[string]string{"name": name},
   380  			},
   381  				types.KubernetesClusterSpecV3{},
   382  			)
   383  			if err != nil {
   384  				return nil, trace.Wrap(err)
   385  			}
   386  			resource, err := types.NewKubernetesServerV3(
   387  				types.Metadata{
   388  					Name: name,
   389  					Labels: map[string]string{
   390  						"label": string(make([]byte, labelSize)),
   391  					},
   392  				},
   393  				types.KubernetesServerSpecV3{
   394  					HostID:  fmt.Sprintf("host-%d", i),
   395  					Cluster: kube,
   396  				},
   397  			)
   398  			if err != nil {
   399  				return nil, trace.Wrap(err)
   400  			}
   401  
   402  			resources = append(resources, any(resource).(T))
   403  		}
   404  	case types.KindWindowsDesktop:
   405  		for i := 0; i < size; i++ {
   406  			var err error
   407  			name := fmt.Sprintf("windows-desktop-%d", i)
   408  			resource, err := types.NewWindowsDesktopV3(
   409  				name,
   410  				map[string]string{"label": string(make([]byte, labelSize))},
   411  				types.WindowsDesktopSpecV3{
   412  					Addr:   "_",
   413  					HostID: "_",
   414  				})
   415  			if err != nil {
   416  				return nil, trace.Wrap(err)
   417  			}
   418  
   419  			resources = append(resources, any(resource).(T))
   420  		}
   421  	case types.KindAppOrSAMLIdPServiceProvider:
   422  		for i := 0; i < size; i++ {
   423  			// Alternate between adding Apps and SAMLIdPServiceProviders. If `i` is even, add an app.
   424  			if i%2 == 0 {
   425  				app, err := types.NewAppV3(types.Metadata{
   426  					Name: fmt.Sprintf("app-%d", i),
   427  				}, types.AppSpecV3{
   428  					URI: "localhost",
   429  				})
   430  				if err != nil {
   431  					return nil, trace.Wrap(err)
   432  				}
   433  
   434  				appServer, err := types.NewAppServerV3(types.Metadata{
   435  					Name: fmt.Sprintf("app-%d", i),
   436  					Labels: map[string]string{
   437  						"label": string(make([]byte, labelSize)),
   438  					},
   439  				}, types.AppServerSpecV3{
   440  					HostID: fmt.Sprintf("host-%d", i),
   441  					App:    app,
   442  				})
   443  				if err != nil {
   444  					return nil, trace.Wrap(err)
   445  				}
   446  
   447  				//nolint:staticcheck // SA1019. TODO(sshah) DELETE IN 17.0
   448  				resource := &types.AppServerOrSAMLIdPServiceProviderV1{
   449  					Resource: &types.AppServerOrSAMLIdPServiceProviderV1_AppServer{
   450  						AppServer: appServer,
   451  					},
   452  				}
   453  
   454  				resources = append(resources, any(resource).(T))
   455  			} else {
   456  				sp := &types.SAMLIdPServiceProviderV1{ResourceHeader: types.ResourceHeader{Metadata: types.Metadata{Name: fmt.Sprintf("saml-app-%d", i), Labels: map[string]string{
   457  					"label": string(make([]byte, labelSize)),
   458  				}}}}
   459  				//nolint:staticcheck // SA1019. TODO(sshah) DELETE IN 17.0
   460  				resource := &types.AppServerOrSAMLIdPServiceProviderV1{
   461  					Resource: &types.AppServerOrSAMLIdPServiceProviderV1_SAMLIdPServiceProvider{
   462  						SAMLIdPServiceProvider: sp,
   463  					},
   464  				}
   465  				resources = append(resources, any(resource).(T))
   466  			}
   467  		}
   468  	case types.KindSAMLIdPServiceProvider:
   469  		for i := 0; i < size; i++ {
   470  			name := fmt.Sprintf("saml-app-%d", i)
   471  			spResource, err := types.NewSAMLIdPServiceProvider(
   472  				types.Metadata{
   473  					Name: name, Labels: map[string]string{
   474  						"label": string(make([]byte, labelSize)),
   475  					},
   476  				},
   477  				types.SAMLIdPServiceProviderSpecV1{
   478  					EntityID: name,
   479  					ACSURL:   name,
   480  				},
   481  			)
   482  			if err != nil {
   483  				return nil, trace.Wrap(err)
   484  			}
   485  
   486  			resources = append(resources, any(spResource).(T))
   487  		}
   488  	default:
   489  		return nil, trace.Errorf("unsupported resource type %s", resourceType)
   490  	}
   491  
   492  	return resources, nil
   493  }
   494  
   495  func TestListResources(t *testing.T) {
   496  	t.Parallel()
   497  	ctx := context.Background()
   498  	srv := startMockServer(t, &listResourcesService{})
   499  
   500  	testCases := map[string]struct {
   501  		resourceType   string
   502  		resourceStruct types.Resource
   503  	}{
   504  		"DatabaseServer": {
   505  			resourceType:   types.KindDatabaseServer,
   506  			resourceStruct: &types.DatabaseServerV3{},
   507  		},
   508  		"ApplicationServer": {
   509  			resourceType:   types.KindAppServer,
   510  			resourceStruct: &types.AppServerV3{},
   511  		},
   512  		"Node": {
   513  			resourceType:   types.KindNode,
   514  			resourceStruct: &types.ServerV2{},
   515  		},
   516  		"KubeServer": {
   517  			resourceType:   types.KindKubeServer,
   518  			resourceStruct: &types.KubernetesServerV3{},
   519  		},
   520  		"WindowsDesktop": {
   521  			resourceType:   types.KindWindowsDesktop,
   522  			resourceStruct: &types.WindowsDesktopV3{},
   523  		},
   524  		"SAMLIdPServiceProvider": {
   525  			resourceType:   types.KindSAMLIdPServiceProvider,
   526  			resourceStruct: &types.SAMLIdPServiceProviderV1{},
   527  		},
   528  	}
   529  
   530  	// Create client
   531  	clt, err := New(ctx, srv.clientCfg())
   532  	require.NoError(t, err)
   533  
   534  	for name, test := range testCases {
   535  		t.Run(name, func(t *testing.T) {
   536  			resp, err := clt.ListResources(ctx, proto.ListResourcesRequest{
   537  				Namespace:    defaults.Namespace,
   538  				Limit:        10,
   539  				ResourceType: test.resourceType,
   540  			})
   541  			require.NoError(t, err)
   542  			require.NotEmpty(t, resp.NextKey)
   543  			require.Len(t, resp.Resources, 10)
   544  			require.IsType(t, test.resourceStruct, resp.Resources[0])
   545  
   546  			// exceed the limit
   547  			_, err = clt.ListResources(ctx, proto.ListResourcesRequest{
   548  				Namespace:    defaults.Namespace,
   549  				Limit:        50,
   550  				ResourceType: test.resourceType,
   551  			})
   552  			require.Error(t, err)
   553  			require.True(t, trace.IsLimitExceeded(err), "trace.IsLimitExceeded failed: err=%v (%T)", err, trace.Unwrap(err))
   554  		})
   555  	}
   556  
   557  	// Test a list with total count returned.
   558  	resp, err := clt.ListResources(ctx, proto.ListResourcesRequest{
   559  		ResourceType:   types.KindNode,
   560  		Limit:          10,
   561  		NeedTotalCount: true,
   562  	})
   563  	require.NoError(t, err)
   564  	require.Equal(t, 50, resp.TotalCount)
   565  }
   566  
   567  func testGetResources[T types.ResourceWithLabels](t *testing.T, clt *Client, kind string) {
   568  	ctx := context.Background()
   569  	expectedResources, err := testResources[T](kind, defaults.Namespace)
   570  	require.NoError(t, err)
   571  
   572  	// Test listing everything at once errors with limit exceeded.
   573  	_, err = clt.ListResources(ctx, proto.ListResourcesRequest{
   574  		Namespace:    defaults.Namespace,
   575  		Limit:        int32(len(expectedResources)),
   576  		ResourceType: kind,
   577  	})
   578  	require.Error(t, err)
   579  	require.True(t, trace.IsLimitExceeded(err), "trace.IsLimitExceeded failed: err=%v (%T)", err, trace.Unwrap(err))
   580  
   581  	// Test getting a page of resources
   582  	page, err := GetResourcePage[T](ctx, clt, &proto.ListResourcesRequest{
   583  		Namespace:      defaults.Namespace,
   584  		ResourceType:   kind,
   585  		NeedTotalCount: true,
   586  	})
   587  	require.NoError(t, err)
   588  	require.Len(t, expectedResources, page.Total)
   589  	require.Empty(t, cmp.Diff(expectedResources[:len(page.Resources)], page.Resources))
   590  
   591  	// Test getting all resources by chunks to handle limit exceeded.
   592  	resources, err := GetAllResources[T](ctx, clt, &proto.ListResourcesRequest{
   593  		Namespace:    defaults.Namespace,
   594  		ResourceType: kind,
   595  	})
   596  	require.NoError(t, err)
   597  	require.Len(t, resources, len(expectedResources))
   598  	require.Empty(t, cmp.Diff(expectedResources, resources))
   599  }
   600  
   601  func TestGetResources(t *testing.T) {
   602  	t.Parallel()
   603  	ctx := context.Background()
   604  	srv := startMockServer(t, &listResourcesService{})
   605  
   606  	// Create client
   607  	clt, err := New(ctx, srv.clientCfg())
   608  	require.NoError(t, err)
   609  
   610  	t.Run("DatabaseServer", func(t *testing.T) {
   611  		t.Parallel()
   612  		testGetResources[types.DatabaseServer](t, clt, types.KindDatabaseServer)
   613  	})
   614  
   615  	t.Run("ApplicationServer", func(t *testing.T) {
   616  		t.Parallel()
   617  		testGetResources[types.AppServer](t, clt, types.KindAppServer)
   618  	})
   619  
   620  	t.Run("Node", func(t *testing.T) {
   621  		t.Parallel()
   622  		testGetResources[types.Server](t, clt, types.KindNode)
   623  	})
   624  
   625  	t.Run("KubeServer", func(t *testing.T) {
   626  		t.Parallel()
   627  		testGetResources[types.KubeServer](t, clt, types.KindKubeServer)
   628  	})
   629  
   630  	t.Run("WindowsDesktop", func(t *testing.T) {
   631  		t.Parallel()
   632  		testGetResources[types.WindowsDesktop](t, clt, types.KindWindowsDesktop)
   633  	})
   634  
   635  	t.Run("AppServerAndSAMLIdPServiceProvider", func(t *testing.T) {
   636  		t.Parallel()
   637  		testGetResources[types.AppServerOrSAMLIdPServiceProvider](t, clt, types.KindAppOrSAMLIdPServiceProvider)
   638  	})
   639  
   640  	t.Run("SAMLIdPServiceProvider", func(t *testing.T) {
   641  		t.Parallel()
   642  		testGetResources[types.SAMLIdPServiceProvider](t, clt, types.KindSAMLIdPServiceProvider)
   643  	})
   644  }
   645  
   646  func TestGetResourcesWithFilters(t *testing.T) {
   647  	t.Parallel()
   648  	ctx := context.Background()
   649  	srv := startMockServer(t, &listResourcesService{})
   650  
   651  	// Create client
   652  	clt, err := New(ctx, srv.clientCfg())
   653  	require.NoError(t, err)
   654  
   655  	testCases := map[string]struct {
   656  		resourceType string
   657  	}{
   658  		"DatabaseServer": {
   659  			resourceType: types.KindDatabaseServer,
   660  		},
   661  		"ApplicationServer": {
   662  			resourceType: types.KindAppServer,
   663  		},
   664  		"Node": {
   665  			resourceType: types.KindNode,
   666  		},
   667  		"KubeServer": {
   668  			resourceType: types.KindKubeServer,
   669  		},
   670  		"WindowsDesktop": {
   671  			resourceType: types.KindWindowsDesktop,
   672  		},
   673  		"AppAndIdPServiceProvider": {
   674  			resourceType: types.KindAppOrSAMLIdPServiceProvider,
   675  		},
   676  		"SAMLIdPServiceProvider": {
   677  			resourceType: types.KindSAMLIdPServiceProvider,
   678  		},
   679  	}
   680  
   681  	for name, test := range testCases {
   682  		name, test := name, test
   683  		t.Run(name, func(t *testing.T) {
   684  			t.Parallel()
   685  			expectedResources, err := testResources[types.ResourceWithLabels](test.resourceType, defaults.Namespace)
   686  			require.NoError(t, err)
   687  
   688  			// Test listing everything at once errors with limit exceeded.
   689  			_, err = clt.ListResources(ctx, proto.ListResourcesRequest{
   690  				Namespace:    defaults.Namespace,
   691  				Limit:        int32(len(expectedResources)),
   692  				ResourceType: test.resourceType,
   693  			})
   694  			require.Error(t, err)
   695  			require.True(t, trace.IsLimitExceeded(err), "trace.IsLimitExceeded failed: err=%v (%T)", err, trace.Unwrap(err))
   696  
   697  			// Test getting all resources by chunks to handle limit exceeded.
   698  			resources, err := GetResourcesWithFilters(ctx, clt, proto.ListResourcesRequest{
   699  				Namespace:    defaults.Namespace,
   700  				ResourceType: test.resourceType,
   701  			})
   702  			require.NoError(t, err)
   703  			require.Len(t, resources, len(expectedResources))
   704  			require.Empty(t, cmp.Diff(expectedResources, resources))
   705  		})
   706  	}
   707  }
   708  
   709  type fakeUnifiedResourcesClient struct {
   710  	resp *proto.ListUnifiedResourcesResponse
   711  	err  error
   712  }
   713  
   714  func (f fakeUnifiedResourcesClient) ListUnifiedResources(ctx context.Context, req *proto.ListUnifiedResourcesRequest) (*proto.ListUnifiedResourcesResponse, error) {
   715  	return f.resp, f.err
   716  }
   717  
   718  // TestGetUnifiedResourcesWithLogins validates that any logins provided
   719  // in a [proto.PaginatedResource] are correctly parsed and applied to
   720  // the corresponding [types.EnrichedResource].
   721  func TestGetUnifiedResourcesWithLogins(t *testing.T) {
   722  	ctx := context.Background()
   723  
   724  	clt := fakeUnifiedResourcesClient{
   725  		resp: &proto.ListUnifiedResourcesResponse{
   726  			Resources: []*proto.PaginatedResource{
   727  				{
   728  					Resource: &proto.PaginatedResource_Node{Node: &types.ServerV2{}},
   729  					Logins:   []string{"alice", "bob"},
   730  				},
   731  				{
   732  					Resource: &proto.PaginatedResource_WindowsDesktop{WindowsDesktop: &types.WindowsDesktopV3{}},
   733  					Logins:   []string{"llama"},
   734  				},
   735  			},
   736  		},
   737  	}
   738  
   739  	resources, _, err := GetUnifiedResourcePage(ctx, clt, &proto.ListUnifiedResourcesRequest{
   740  		SortBy: types.SortBy{
   741  			IsDesc: false,
   742  			Field:  types.ResourceSpecHostname,
   743  		},
   744  		IncludeLogins: true,
   745  	})
   746  	require.NoError(t, err)
   747  
   748  	require.Len(t, resources, len(clt.resp.Resources))
   749  
   750  	for _, enriched := range resources {
   751  		switch enriched.ResourceWithLabels.(type) {
   752  		case *types.ServerV2:
   753  			assert.Equal(t, enriched.Logins, clt.resp.Resources[0].Logins)
   754  		case *types.WindowsDesktopV3:
   755  			assert.Equal(t, enriched.Logins, clt.resp.Resources[1].Logins)
   756  		}
   757  	}
   758  }