From b412b0d80fbb4c378bc8330fa5d1fd283e47b05a Mon Sep 17 00:00:00 2001 From: Saketh Somaraju Date: Mon, 2 Feb 2026 15:29:31 +0530 Subject: [PATCH 1/2] set explicit limit when listing organizations --- cloud/organization/organization.go | 5 +++- cloud/organization/organization_test.go | 33 +++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/cloud/organization/organization.go b/cloud/organization/organization.go index cef5f37ad..7ded3f1e8 100644 --- a/cloud/organization/organization.go +++ b/cloud/organization/organization.go @@ -42,7 +42,10 @@ func newTableOut() *printutil.Table { } func ListOrganizations(platformCoreClient astroplatformcore.CoreClient) ([]astroplatformcore.Organization, error) { - organizationListParams := &astroplatformcore.ListOrganizationsParams{} + limit := 1000 + organizationListParams := &astroplatformcore.ListOrganizationsParams{ + Limit: &limit, + } resp, err := platformCoreClient.ListOrganizationsWithResponse(http_context.Background(), organizationListParams) if err != nil { return nil, err diff --git a/cloud/organization/organization_test.go b/cloud/organization/organization_test.go index 9eac76c72..207fa5fee 100644 --- a/cloud/organization/organization_test.go +++ b/cloud/organization/organization_test.go @@ -15,6 +15,7 @@ import ( astroplatformcore_mocks "github.com/astronomer/astro-cli/astro-client-platform-core/mocks" "github.com/astronomer/astro-cli/config" testUtil "github.com/astronomer/astro-cli/pkg/testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" "github.com/stretchr/testify/suite" ) @@ -115,6 +116,38 @@ func (s *Suite) TestList() { }) } +func TestListOrganizations(t *testing.T) { + mockClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) + + limit := 1000 + orgs := []astroplatformcore.Organization{ + {Id: "org-1", Name: "Org 1"}, + {Id: "org-2", Name: "Org 2"}, + } + + resp := &astroplatformcore.ListOrganizationsResponse{ + HTTPResponse: &http.Response{StatusCode: http.StatusOK}, + JSON200: &astroplatformcore.OrganizationsPaginated{ + Organizations: orgs, + }, + } + + mockClient.On("ListOrganizationsWithResponse", mock.Anything, mock.MatchedBy(func(p *astroplatformcore.ListOrganizationsParams) bool { + return p != nil && p.Limit != nil && *p.Limit == limit + }), + mock.Anything, + ). + Return(resp, nil). + Once() + + organizations, err := ListOrganizations(mockClient) + assert.NoError(t, err) + assert.Len(t, organizations, 2) + assert.Equal(t, orgs, organizations) + + mockClient.AssertExpectations(t) +} + func (s *Suite) TestGetOrganizationSelection() { // initialize empty config testUtil.InitTestConfig(testUtil.LocalPlatform) From 82f0c07b89f66ca2a3c385aefbce9feb255fb536 Mon Sep 17 00:00:00 2001 From: Saketh Somaraju Date: Mon, 16 Feb 2026 08:36:03 +0530 Subject: [PATCH 2/2] fix tests --- cloud/organization/organization.go | 3 +- cloud/organization/organization_test.go | 42 ++++++++++++++----------- cmd/cloud/setup_test.go | 24 ++++++++++---- 3 files changed, 44 insertions(+), 25 deletions(-) diff --git a/cloud/organization/organization.go b/cloud/organization/organization.go index 7ded3f1e8..b1ce36996 100644 --- a/cloud/organization/organization.go +++ b/cloud/organization/organization.go @@ -21,6 +21,7 @@ import ( const ( AstronomerConnectionErrMsg = "cannot connect to Astronomer. Try to log in with astro login or check your internet connection and user permissions. If you are using an API Key or Token make sure your context is correct.\n\nDetails" + listOrganizationsLimit = 1000 ) var ( @@ -42,7 +43,7 @@ func newTableOut() *printutil.Table { } func ListOrganizations(platformCoreClient astroplatformcore.CoreClient) ([]astroplatformcore.Organization, error) { - limit := 1000 + limit := listOrganizationsLimit organizationListParams := &astroplatformcore.ListOrganizationsParams{ Limit: &limit, } diff --git a/cloud/organization/organization_test.go b/cloud/organization/organization_test.go index 207fa5fee..370a96a02 100644 --- a/cloud/organization/organization_test.go +++ b/cloud/organization/organization_test.go @@ -83,13 +83,19 @@ var ( errNetwork = errors.New("network error") ) +func matchListOrganizationsLimit() interface{} { + return mock.MatchedBy(func(p *astroplatformcore.ListOrganizationsParams) bool { + return p != nil && p.Limit != nil && *p.Limit == listOrganizationsLimit + }) +} + func (s *Suite) TestList() { // initialize empty config testUtil.InitTestConfig(testUtil.LocalPlatform) s.Run("organization list success", func() { mockClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() buf := new(bytes.Buffer) err := List(buf, mockClient) @@ -99,7 +105,7 @@ func (s *Suite) TestList() { s.Run("organization network error", func() { mockClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(nil, errNetwork).Once() + mockClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(nil, errNetwork).Once() buf := new(bytes.Buffer) err := List(buf, mockClient) s.Contains(err.Error(), "network error") @@ -108,7 +114,7 @@ func (s *Suite) TestList() { s.Run("organization list error", func() { mockClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockErrorResponse, nil).Once() + mockClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockErrorResponse, nil).Once() buf := new(bytes.Buffer) err := List(buf, mockClient) s.Contains(err.Error(), "failed to fetch organizations") @@ -119,7 +125,7 @@ func (s *Suite) TestList() { func TestListOrganizations(t *testing.T) { mockClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - limit := 1000 + limit := listOrganizationsLimit orgs := []astroplatformcore.Organization{ {Id: "org-1", Name: "Org 1"}, {Id: "org-2", Name: "Org 2"}, @@ -154,7 +160,7 @@ func (s *Suite) TestGetOrganizationSelection() { s.Run("get organiation selection success", func() { mockClient := new(astrocore_mocks.ClientWithResponsesInterface) mockPlatformCoreClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() // mock os.Stdin input := []byte("1") @@ -177,7 +183,7 @@ func (s *Suite) TestGetOrganizationSelection() { s.Run("get organization selection list error", func() { mockClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockErrorResponse, nil).Once() + mockClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockErrorResponse, nil).Once() buf := new(bytes.Buffer) _, err := getOrganizationSelection(buf, mockClient) @@ -187,7 +193,7 @@ func (s *Suite) TestGetOrganizationSelection() { s.Run("get organization selection select error", func() { mockClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() // mock os.Stdin input := []byte("3") @@ -214,7 +220,7 @@ func (s *Suite) TestSwitch() { s.Run("successful switch with name", func() { mockCoreClient := new(astrocore_mocks.ClientWithResponsesInterface) mockPlatformCoreClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() CheckUserSession = func(c *config.Context, coreClient astrocore.CoreClient, platformCoreClient astroplatformcore.CoreClient, out io.Writer) error { return nil } @@ -229,7 +235,7 @@ func (s *Suite) TestSwitch() { s.Run("switching to a current organization", func() { mockCoreClient := new(astrocore_mocks.ClientWithResponsesInterface) mockPlatformCoreClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() buf := new(bytes.Buffer) err := Switch("org1", mockCoreClient, mockPlatformCoreClient, buf, false) @@ -242,7 +248,7 @@ func (s *Suite) TestSwitch() { s.Run("successful switch without name", func() { mockCoreClient := new(astrocore_mocks.ClientWithResponsesInterface) mockPlatformCoreClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() CheckUserSession = func(c *config.Context, coreClient astrocore.CoreClient, platformCoreClient astroplatformcore.CoreClient, out io.Writer) error { return nil } @@ -267,7 +273,7 @@ func (s *Suite) TestSwitch() { s.Run("failed switch wrong name", func() { mockCoreClient := new(astrocore_mocks.ClientWithResponsesInterface) mockPlatformCoreClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() CheckUserSession = func(c *config.Context, coreClient astrocore.CoreClient, platformCoreClient astroplatformcore.CoreClient, out io.Writer) error { return nil } @@ -281,7 +287,7 @@ func (s *Suite) TestSwitch() { s.Run("failed switch bad selection", func() { mockCoreClient := new(astrocore_mocks.ClientWithResponsesInterface) mockPlatformCoreClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() CheckUserSession = func(c *config.Context, coreClient astrocore.CoreClient, platformCoreClient astroplatformcore.CoreClient, out io.Writer) error { return nil } @@ -316,7 +322,7 @@ func (s *Suite) TestSwitch() { } mockCoreClient := new(astrocore_mocks.ClientWithResponsesInterface) mockPlatformCoreClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() CheckUserSession = func(c *config.Context, coreClient astrocore.CoreClient, platformCoreClient astroplatformcore.CoreClient, out io.Writer) error { return nil } @@ -392,7 +398,7 @@ func (s *Suite) TestExportAuditLogs() { s.Run("export audit logs success", func() { mockPlatformClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) mockClient := new(astrocore_mocks.ClientWithResponsesInterface) - mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() mockClient.On("GetOrganizationAuditLogsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockOKAuditLogResponse, nil).Once() err := ExportAuditLogs(mockClient, mockPlatformClient, "", "", 1) s.NoError(err) @@ -402,7 +408,7 @@ func (s *Suite) TestExportAuditLogs() { s.Run("export audit logs and select org success", func() { mockPlatformClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) mockClient := new(astrocore_mocks.ClientWithResponsesInterface) - mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() mockClient.On("GetOrganizationAuditLogsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockOKAuditLogResponse, nil).Once() err := ExportAuditLogs(mockClient, mockPlatformClient, "org1", "", 1) s.NoError(err) @@ -412,7 +418,7 @@ func (s *Suite) TestExportAuditLogs() { s.Run("export failure", func() { mockPlatformClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) mockClient := new(astrocore_mocks.ClientWithResponsesInterface) - mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() mockClient.On("GetOrganizationAuditLogsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(nil, errNetwork).Once() err := ExportAuditLogs(mockClient, mockPlatformClient, "", "", 1) s.Contains(err.Error(), "network error") @@ -422,7 +428,7 @@ func (s *Suite) TestExportAuditLogs() { s.Run("list failure", func() { mockPlatformClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) mockClient := new(astrocore_mocks.ClientWithResponsesInterface) - mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(nil, errNetwork).Once() + mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(nil, errNetwork).Once() err := ExportAuditLogs(mockClient, mockPlatformClient, "org1", "", 1) s.Contains(err.Error(), "network error") mockPlatformClient.AssertExpectations(s.T()) @@ -431,7 +437,7 @@ func (s *Suite) TestExportAuditLogs() { s.Run("organization list error", func() { mockPlatformClient := new(astroplatformcore_mocks.ClientWithResponsesInterface) mockClient := new(astrocore_mocks.ClientWithResponsesInterface) - mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOKResponse, nil).Once() + mockPlatformClient.On("ListOrganizationsWithResponse", mock.Anything, matchListOrganizationsLimit()).Return(&mockOKResponse, nil).Once() mockClient.On("GetOrganizationAuditLogsWithResponse", mock.Anything, mock.Anything, mock.Anything).Return(&mockOKAuditLogResponseError, nil).Once() err := ExportAuditLogs(mockClient, mockPlatformClient, "", "", 1) s.Contains(err.Error(), "failed to fetch organizations audit logs") diff --git a/cmd/cloud/setup_test.go b/cmd/cloud/setup_test.go index 0f3974cc4..152adf711 100644 --- a/cmd/cloud/setup_test.go +++ b/cmd/cloud/setup_test.go @@ -472,7 +472,9 @@ func TestCheckAPIToken(t *testing.T) { return &mockClaims, nil } - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOrgsResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, mock.MatchedBy(func(p *astroplatformcore.ListOrganizationsParams) bool { + return p != nil && p.Limit != nil && *p.Limit == 1000 + })).Return(&mockOrgsResponse, nil).Once() t.Setenv("ASTRO_API_TOKEN", "token") @@ -495,7 +497,9 @@ func TestCheckAPIToken(t *testing.T) { return nil, errors.New("Failed to parse token") } - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOrgsResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, mock.MatchedBy(func(p *astroplatformcore.ListOrganizationsParams) bool { + return p != nil && p.Limit != nil && *p.Limit == 1000 + })).Return(&mockOrgsResponse, nil).Once() t.Setenv("ASTRO_API_TOKEN", "token") @@ -517,7 +521,9 @@ func TestCheckAPIToken(t *testing.T) { return &mockClaims, nil } - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOrgsResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, mock.MatchedBy(func(p *astroplatformcore.ListOrganizationsParams) bool { + return p != nil && p.Limit != nil && *p.Limit == 1000 + })).Return(&mockOrgsResponse, nil).Once() t.Setenv("ASTRO_API_TOKEN", "token") err := config.ResetCurrentContext() @@ -551,7 +557,9 @@ func TestCheckAPIToken(t *testing.T) { return &mockClaims, nil } - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOrgsResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, mock.MatchedBy(func(p *astroplatformcore.ListOrganizationsParams) bool { + return p != nil && p.Limit != nil && *p.Limit == 1000 + })).Return(&mockOrgsResponse, nil).Once() t.Setenv("ASTRO_API_TOKEN", "token") @@ -591,7 +599,9 @@ func TestCheckAPIToken(t *testing.T) { return &mockClaims, nil } - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOrgsResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, mock.MatchedBy(func(p *astroplatformcore.ListOrganizationsParams) bool { + return p != nil && p.Limit != nil && *p.Limit == 1000 + })).Return(&mockOrgsResponse, nil).Once() t.Setenv("ASTRO_API_TOKEN", "token") @@ -630,7 +640,9 @@ func TestCheckAPIToken(t *testing.T) { return &mockClaims, nil } - mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, &astroplatformcore.ListOrganizationsParams{}).Return(&mockOrgsResponse, nil).Once() + mockPlatformCoreClient.On("ListOrganizationsWithResponse", mock.Anything, mock.MatchedBy(func(p *astroplatformcore.ListOrganizationsParams) bool { + return p != nil && p.Limit != nil && *p.Limit == 1000 + })).Return(&mockOrgsResponse, nil).Once() t.Setenv("ASTRO_API_TOKEN", "token")