mirror of
https://github.com/fleetdm/fleet
synced 2026-04-24 06:57:21 +00:00
115 lines
3.4 KiB
Go
115 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"strconv"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/fleetdm/fleet/v4/server/datastore/cached_mysql"
|
|
"github.com/fleetdm/fleet/v4/server/fleet"
|
|
"github.com/fleetdm/fleet/v4/server/mock"
|
|
"github.com/fleetdm/fleet/v4/server/service"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/urfave/cli/v2"
|
|
)
|
|
|
|
// runServerWithMockedDS runs the fleet server with several mocked DS methods.
|
|
//
|
|
// NOTE: Assumes the current session is always from the admin user (see ds.SessionByKeyFunc below).
|
|
func runServerWithMockedDS(t *testing.T, opts ...*service.TestServerOpts) (*httptest.Server, *mock.Store) {
|
|
ds := new(mock.Store)
|
|
var users []*fleet.User
|
|
var admin *fleet.User
|
|
ds.NewUserFunc = func(ctx context.Context, user *fleet.User) (*fleet.User, error) {
|
|
if user.GlobalRole != nil && *user.GlobalRole == fleet.RoleAdmin {
|
|
admin = user
|
|
}
|
|
users = append(users, user)
|
|
return user, nil
|
|
}
|
|
ds.SessionByKeyFunc = func(ctx context.Context, key string) (*fleet.Session, error) {
|
|
return &fleet.Session{
|
|
CreateTimestamp: fleet.CreateTimestamp{CreatedAt: time.Now()},
|
|
ID: 1,
|
|
AccessedAt: time.Now(),
|
|
UserID: admin.ID,
|
|
Key: key,
|
|
}, nil
|
|
}
|
|
ds.MarkSessionAccessedFunc = func(ctx context.Context, session *fleet.Session) error {
|
|
return nil
|
|
}
|
|
ds.UserByIDFunc = func(ctx context.Context, id uint) (*fleet.User, error) {
|
|
return admin, nil
|
|
}
|
|
ds.ListUsersFunc = func(ctx context.Context, opt fleet.UserListOptions) ([]*fleet.User, error) {
|
|
return users, nil
|
|
}
|
|
ds.AppConfigFunc = func(ctx context.Context) (*fleet.AppConfig, error) {
|
|
return &fleet.AppConfig{}, nil
|
|
}
|
|
|
|
cachedDS := cached_mysql.New(ds)
|
|
_, server := service.RunServerForTestsWithDS(t, cachedDS, opts...)
|
|
os.Setenv("FLEET_SERVER_ADDRESS", server.URL)
|
|
|
|
return server, ds
|
|
}
|
|
|
|
func runAppForTest(t *testing.T, args []string) string {
|
|
w, err := runAppNoChecks(args)
|
|
require.NoError(t, err)
|
|
return w.String()
|
|
}
|
|
|
|
func runAppCheckErr(t *testing.T, args []string, errorMsg string) string {
|
|
w, err := runAppNoChecks(args)
|
|
require.Error(t, err)
|
|
require.Equal(t, errorMsg, err.Error())
|
|
return w.String()
|
|
}
|
|
|
|
func runAppNoChecks(args []string) (*bytes.Buffer, error) {
|
|
// first arg must be the binary name. Allow tests to omit it.
|
|
args = append([]string{""}, args...)
|
|
|
|
w := new(bytes.Buffer)
|
|
app := createApp(nil, w, os.Stderr, noopExitErrHandler)
|
|
err := app.Run(args)
|
|
return w, err
|
|
}
|
|
|
|
func runWithErrWriter(args []string, errWriter io.Writer) (*bytes.Buffer, error) {
|
|
args = append([]string{""}, args...)
|
|
|
|
w := new(bytes.Buffer)
|
|
app := createApp(nil, w, errWriter, noopExitErrHandler)
|
|
err := app.Run(args)
|
|
return w, err
|
|
}
|
|
|
|
func noopExitErrHandler(c *cli.Context, err error) {}
|
|
|
|
func serveMDMBootstrapPackage(t *testing.T, pkgPath, pkgName string) (*httptest.Server, int) {
|
|
pkgBytes, err := os.ReadFile(pkgPath)
|
|
require.NoError(t, err)
|
|
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Length", strconv.Itoa(len(pkgBytes)))
|
|
w.Header().Set("Content-Type", "application/octet-stream")
|
|
w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment;filename="%s"`, pkgName))
|
|
if n, err := w.Write(pkgBytes); err != nil {
|
|
require.NoError(t, err)
|
|
require.Equal(t, len(pkgBytes), n)
|
|
}
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
return srv, len(pkgBytes)
|
|
}
|