Skip to content

Commit

Permalink
Refactor new command (#4153)
Browse files Browse the repository at this point in the history 
* Refactor new command and templates

Signed-off-by: Daniel González Lopes <[email protected]>

* Move Cloud to another command

Signed-off-by: Daniel González Lopes <[email protected]>

* Fix tests

Signed-off-by: Daniel González Lopes <[email protected]>

* Small changes to the templates

Signed-off-by: Daniel González Lopes <[email protected]>

* Move everything to one command based on feedback

Signed-off-by: Daniel González Lopes <[email protected]>

* Fix help text

Signed-off-by: Daniel González Lopes <[email protected]>

* Simplify cloud logic

Signed-off-by: Daniel González Lopes <[email protected]>

* Make linter happy

Signed-off-by: Daniel González Lopes <[email protected]>

* Fix package name and description

Signed-off-by: Daniel González Lopes <[email protected]>

* Use Authorization header instead of User-ID

Signed-off-by: Daniel González Lopes <[email protected]>

* Remove extra vus

Signed-off-by: Daniel González Lopes <[email protected]>

* Improve example text based on feedback

Signed-off-by: Daniel González Lopes <[email protected]>

* Add new lines in the browser template based on feedback

Signed-off-by: Daniel González Lopes <[email protected]>

* Add vus/iterations to the Browser example

Signed-off-by: Daniel González Lopes <[email protected]>

* Add catch + fail in the Browser template

Signed-off-by: Daniel González Lopes <[email protected]>

* Rename templates folder

Signed-off-by: Daniel González Lopes <[email protected]>

* Fix copy of force command

Signed-off-by: Daniel González Lopes <[email protected]>

* Remove the panic

Signed-off-by: Daniel González Lopes <[email protected]>

* Add constants for template types

Signed-off-by: Daniel González Lopes <[email protected]>

* Use exec.test.abort instead of throw

Signed-off-by: Daniel González Lopes <[email protected]>

* Make linter happy

Signed-off-by: Daniel González Lopes <[email protected]>

* Add package comment to templates.go

---------

Signed-off-by: Daniel González Lopes <[email protected]>
  • Loading branch information
@dgzlopes
dgzlopes authored Jan 22, 2025
1 parent a9854f1 commit 6c4752c
Show file tree
Hide file tree
Showing 6 changed files with 294 additions and 103 deletions.
158 changes: 57 additions & 101 deletions cmd/new.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,99 +2,33 @@ package cmd

import (
"fmt"
"path"
"text/template"

"github.com/fatih/color"
"github.com/spf13/cobra"
"github.com/spf13/pflag"
"go.k6.io/k6/cmd/state"
"go.k6.io/k6/cmd/templates"
"go.k6.io/k6/lib/fsext"
)

const defaultNewScriptName = "script.js"

//nolint:gochecknoglobals
var defaultNewScriptTemplate = template.Must(template.New("new").Parse(`import http from 'k6/http';
import { sleep } from 'k6';
export const options = {
// A number specifying the number of VUs to run concurrently.
vus: 10,
// A string specifying the total duration of the test run.
duration: '30s',
// The following section contains configuration options for execution of this
// test script in Grafana Cloud.
//
// See https://grafana.com/docs/grafana-cloud/k6/get-started/run-cloud-tests-from-the-cli/
// to learn about authoring and running k6 test scripts in Grafana k6 Cloud.
//
// cloud: {
// // The ID of the project to which the test is assigned in the k6 Cloud UI.
// // By default tests are executed in default project.
// projectID: "",
// // The name of the test in the k6 Cloud UI.
// // Test runs with the same name will be grouped.
// name: "{{ .ScriptName }}"
// },
// Uncomment this section to enable the use of Browser API in your tests.
//
// See https://grafana.com/docs/k6/latest/using-k6-browser/running-browser-tests/ to learn more
// about using Browser API in your test scripts.
//
// scenarios: {
// // The scenario name appears in the result summary, tags, and so on.
// // You can give the scenario any name, as long as each name in the script is unique.
// ui: {
// // Executor is a mandatory parameter for browser-based tests.
// // Shared iterations in this case tells k6 to reuse VUs to execute iterations.
// //
// // See https://grafana.com/docs/k6/latest/using-k6/scenarios/executors/ for other executor types.
// executor: 'shared-iterations',
// options: {
// browser: {
// // This is a mandatory parameter that instructs k6 to launch and
// // connect to a chromium-based browser, and use it to run UI-based
// // tests.
// type: 'chromium',
// },
// },
// },
// }
};
// The function that defines VU logic.
//
// See https://grafana.com/docs/k6/latest/examples/get-started-with-k6/ to learn more
// about authoring k6 scripts.
//
export default function() {
http.get('https://test.k6.io');
sleep(1);
}
`))

type initScriptTemplateArgs struct {
ScriptName string
}

// newScriptCmd represents the `k6 new` command
type newScriptCmd struct {
gs *state.GlobalState
overwriteFiles bool
templateType string
projectID string
}

func (c *newScriptCmd) flagSet() *pflag.FlagSet {
flags := pflag.NewFlagSet("", pflag.ContinueOnError)
flags.SortFlags = false
flags.BoolVarP(&c.overwriteFiles, "force", "f", false, "Overwrite existing files")

flags.BoolVarP(&c.overwriteFiles, "force", "f", false, "overwrite existing files")
flags.StringVar(&c.templateType, "template", "minimal", "template type (choices: minimal, protocol, browser)")
flags.StringVar(&c.projectID, "project-id", "", "specify the Grafana Cloud project ID for the test")
return flags
}

func (c *newScriptCmd) run(cmd *cobra.Command, args []string) error { //nolint:revive
func (c *newScriptCmd) run(_ *cobra.Command, args []string) error {
target := defaultNewScriptName
if len(args) > 0 {
target = args[0]
Expand All @@ -104,64 +38,86 @@ func (c *newScriptCmd) run(cmd *cobra.Command, args []string) error { //nolint:r
if err != nil {
return err
}

if fileExists && !c.overwriteFiles {
return fmt.Errorf("%s already exists, please use the `--force` flag if you want overwrite it", target)
return fmt.Errorf("%s already exists. Use the `--force` flag to overwrite it", target)
}

fd, err := c.gs.FS.Create(target)
if err != nil {
return err
}

var closeErr error
defer func() {
_ = fd.Close() // we may think to check the error and log
if cerr := fd.Close(); cerr != nil {
if _, err := fmt.Fprintf(c.gs.Stderr, "error closing file: %v\n", cerr); err != nil {
closeErr = fmt.Errorf("error writing error message to stderr: %w", err)
} else {
closeErr = cerr
}
}
}()

if err := defaultNewScriptTemplate.Execute(fd, initScriptTemplateArgs{
ScriptName: path.Base(target),
}); err != nil {
if closeErr != nil {
return closeErr
}

tm, err := templates.NewTemplateManager()
if err != nil {
return fmt.Errorf("error initializing template manager: %w", err)
}

tmpl, err := tm.GetTemplate(c.templateType)
if err != nil {
return fmt.Errorf("error retrieving template: %w", err)
}

argsStruct := templates.TemplateArgs{
ScriptName: target,
ProjectID: c.projectID,
}

if err := templates.ExecuteTemplate(fd, tmpl, argsStruct); err != nil {
return err
}

valueColor := getColor(c.gs.Flags.NoColor || !c.gs.Stdout.IsTTY, color.Bold)
printToStdout(c.gs, fmt.Sprintf(
"Initialized a new k6 test script in %s. You can now execute it by running `%s run %s`.\n",
valueColor.Sprint(target),
c.gs.BinaryName,
target,
))
if _, err := fmt.Fprintf(c.gs.Stdout, "New script created: %s (%s template).\n", target, c.templateType); err != nil {
return err
}

return nil
}

func getCmdNewScript(gs *state.GlobalState) *cobra.Command {
c := &newScriptCmd{gs: gs}

exampleText := getExampleText(gs, `
# Create a minimal k6 script in the current directory. By default, k6 creates script.js.
{{.}} new
exampleText := getExampleText(c.gs, `
# Create a new k6 script with the default template
$ {{.}} new
# Specify a file name when creating a script
$ {{.}} new test.js
# Overwrite an existing file
$ {{.}} new -f test.js
# Create a minimal k6 script in the current directory and store it in test.js
{{.}} new test.js
# Create a script using a specific template
$ {{.}} new --template protocol
# Overwrite existing test.js with a minimal k6 script
{{.}} new -f test.js`[1:])
# Create a cloud-ready script with a specific project ID
$ {{.}} new --project-id 12315`[1:])

initCmd := &cobra.Command{
Use: "new",
Use: "new [file]",
Short: "Create and initialize a new k6 script",
Long: `Create and initialize a new k6 script.
Long: `Create and initialize a new k6 script using one of the predefined templates.
This command will create a minimal k6 script in the current directory and
store it in the file specified by the first argument. If no argument is
provided, the script will be stored in script.js.
This command will not overwrite existing files.`,
By default, the script will be named script.js unless a different name is specified.`,
Example: exampleText,
Args: cobra.MaximumNArgs(1),
RunE: c.run,
}
initCmd.Flags().AddFlagSet(c.flagSet())

initCmd.Flags().AddFlagSet(c.flagSet())
return initCmd
}
28 changes: 26 additions & 2 deletions cmd/new_test.go
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package cmd

import (
"fmt"
"testing"

"github.com/stretchr/testify/assert"
Expand Down Expand Up @@ -56,7 +55,6 @@ func TestNewScriptCmd(t *testing.T) {

jsData := string(data)
assert.Contains(t, jsData, "export const options = {")
assert.Contains(t, jsData, fmt.Sprintf(`name: "%s"`, testCase.expectedCloudName))
assert.Contains(t, jsData, "export default function() {")
})
}
Expand Down Expand Up @@ -95,3 +93,29 @@ func TestNewScriptCmd_FileExists_Overwrite(t *testing.T) {
assert.Contains(t, string(data), "export const options = {")
assert.Contains(t, string(data), "export default function() {")
}

func TestNewScriptCmd_InvalidTemplateType(t *testing.T) {
t.Parallel()

ts := tests.NewGlobalTestState(t)
ts.CmdArgs = []string{"k6", "new", "--template", "invalid-template"}

ts.ExpectedExitCode = -1

newRootCommand(ts.GlobalState).execute()
assert.Contains(t, ts.Stderr.String(), "invalid template type")
}

func TestNewScriptCmd_ProjectID(t *testing.T) {
t.Parallel()

ts := tests.NewGlobalTestState(t)
ts.CmdArgs = []string{"k6", "new", "--project-id", "1422"}

newRootCommand(ts.GlobalState).execute()

data, err := fsext.ReadFile(ts.FS, defaultNewScriptName)
require.NoError(t, err)

assert.Contains(t, string(data), "projectID: 1422")
}
63 changes: 63 additions & 0 deletions cmd/templates/browser.js
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
import http from "k6/http";
import exec from 'k6/execution';
import { browser } from "k6/browser";
import { sleep, check, fail } from 'k6';

const BASE_URL = __ENV.BASE_URL || "https://quickpizza.grafana.com";

export const options = {
scenarios: {
ui: {
executor: "shared-iterations",
vus: 1,
iterations: 1,
options: {
browser: {
type: "chromium",
},
},
},
},{{ if .ProjectID }}
cloud: {
projectID: {{ .ProjectID }},
name: "{{ .ScriptName }}",
},{{ end }}
};

export function setup() {
let res = http.get(BASE_URL);
if (res.status !== 200) {
exec.test.abort(`Got unexpected status code ${res.status} when trying to setup. Exiting.`);
}
}

export default async function() {
let checkData;
const page = await browser.newPage();

try {
await page.goto(BASE_URL);

checkData = await page.locator("h1").textContent();
check(page, {
header: checkData === "Looking to break out of your pizza routine?",
});

await page.locator('//button[. = "Pizza, Please!"]').click();
await page.waitForTimeout(500);

await page.screenshot({ path: "screenshot.png" });

checkData = await page.locator("div#recommendations").textContent();
check(page, {
recommendation: checkData !== "",
});
} catch (error) {
fail(`Browser iteration failed: ${error.message}`);
} finally {
await page.close();
}

sleep(1);
}

17 changes: 17 additions & 0 deletions cmd/templates/minimal.js
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,17 @@
import http from 'k6/http';
import { sleep, check } from 'k6';

export const options = {
vus: 10,
duration: '30s',{{ if .ProjectID }}
cloud: {
projectID: {{ .ProjectID }},
name: "{{ .ScriptName }}",
},{{ end }}
};

export default function() {
let res = http.get('https://quickpizza.grafana.com');
check(res, { "status is 200": (res) => res.status === 200 });
sleep(1);
}
Loading

0 comments on commit 6c4752c

Please sign in to comment.