Raise correct error when a resource doesn't exist

internal/cmd/branch/create.go

@@ -86,10 +86,18 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 			if err != nil {
 				switch cmdutil.ErrCode(err) {
 				case ps.ErrNotFound:
+					if _, orgErr := client.Organizations.Get(ctx, &ps.GetOrganizationRequest{
+						Organization: ch.Config.Organization,
+					}); cmdutil.ErrCode(orgErr) == ps.ErrNotFound {
+						return cmdutil.HandleNotFoundWithServiceTokenCheck(
+							ctx, cmd, ch.Config, ch.Client, err, "read_branch",
+							"organization %s does not exist",
+							printer.BoldBlue(ch.Config.Organization))
+					}
 					return cmdutil.HandleNotFoundWithServiceTokenCheck(
 						ctx, cmd, ch.Config, ch.Client, err, "read_branch",
-						"branch %s does not exist in database %s (organization: %s)",
-						printer.BoldBlue(branch), printer.BoldBlue(source), printer.BoldBlue(ch.Config.Organization))
+						"source database %s does not exist in organization %s",
+						printer.BoldBlue(source), printer.BoldBlue(ch.Config.Organization))
 				default:
 					return cmdutil.HandleError(err)
 				}

internal/cmd/branch/create_test.go

@@ -529,3 +529,96 @@ func TestBranch_CreateCmd_GenuineNotFound(t *testing.T) {
 	c.Assert(branchSvc.CreateFnInvoked, qt.IsTrue)
 	c.Assert(orgSvc.ListFnInvoked, qt.IsTrue)
 }
+
+func TestBranch_CreateCmd_OrgNotFound(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	org := "nonexistent-org"
+	db := "planetscale"
+	branch := "development"
+
+	orgSvc := &mock.OrganizationsService{
+		GetFn: func(ctx context.Context, req *ps.GetOrganizationRequest) (*ps.Organization, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			return nil, &ps.Error{Code: ps.ErrNotFound}
+		},
+	}
+
+	dbSvc := &mock.DatabaseService{
+		GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) {
+			return nil, &ps.Error{Code: ps.ErrNotFound}
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Organizations: orgSvc,
+				Databases:     dbSvc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, branch})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err.Error(), qt.Contains, "organization "+org+" does not exist")
+	c.Assert(err.Error(), qt.Not(qt.Contains), branch)
+	c.Assert(orgSvc.GetFnInvoked, qt.IsTrue)
+}
+
+func TestBranch_CreateCmd_SourceDatabaseNotFound(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	org := "planetscale"
+	db := "nonexistent-db"
+	branch := "development"
+
+	orgSvc := &mock.OrganizationsService{
+		GetFn: func(ctx context.Context, req *ps.GetOrganizationRequest) (*ps.Organization, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			return &ps.Organization{Name: org}, nil
+		},
+	}
+
+	dbSvc := &mock.DatabaseService{
+		GetFn: func(ctx context.Context, req *ps.GetDatabaseRequest) (*ps.Database, error) {
+			c.Assert(req.Database, qt.Equals, db)
+			return nil, &ps.Error{Code: ps.ErrNotFound}
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: org},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Organizations: orgSvc,
+				Databases:     dbSvc,
+			}, nil
+		},
+	}
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{db, branch})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNotNil)
+	c.Assert(err.Error(), qt.Contains, "source database "+db+" does not exist in organization "+org)
+	c.Assert(err.Error(), qt.Not(qt.Contains), branch)
+	c.Assert(orgSvc.GetFnInvoked, qt.IsTrue)
+}