From 73a08f0045287942cad4093a97503de642647ee1 Mon Sep 17 00:00:00 2001
From: Dave Henderson <dhenderson@gmail.com>
Date: Sun, 5 Aug 2018 05:03:18 -0400
Subject: [PATCH] promtool - Adding --step flag to 'query range' subcommand
 (#4454)

Signed-off-by: Dave Henderson <dhenderson@gmail.com>
---
 cmd/promtool/main.go      | 13 ++++---
 cmd/promtool/main_test.go | 75 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 5 deletions(-)
 create mode 100644 cmd/promtool/main_test.go

diff --git a/cmd/promtool/main.go b/cmd/promtool/main.go
index d5daca383..88a601a01 100644
--- a/cmd/promtool/main.go
+++ b/cmd/promtool/main.go
@@ -74,6 +74,7 @@ func main() {
 	queryRangeExpr := queryRangeCmd.Arg("expr", "PromQL query expression.").Required().String()
 	queryRangeBegin := queryRangeCmd.Flag("start", "Query range start time (RFC3339 or Unix timestamp).").String()
 	queryRangeEnd := queryRangeCmd.Flag("end", "Query range end time (RFC3339 or Unix timestamp).").String()
+	queryRangeStep := queryRangeCmd.Flag("step", "Query step size (duration).").Duration()
 
 	querySeriesCmd := queryCmd.Command("series", "Run series query.")
 	querySeriesServer := querySeriesCmd.Arg("server", "Prometheus server to query.").Required().URL()
@@ -110,7 +111,7 @@ func main() {
 		os.Exit(QueryInstant(*queryServer, *queryExpr))
 
 	case queryRangeCmd.FullCommand():
-		os.Exit(QueryRange(*queryRangeServer, *queryRangeExpr, *queryRangeBegin, *queryRangeEnd))
+		os.Exit(QueryRange(*queryRangeServer, *queryRangeExpr, *queryRangeBegin, *queryRangeEnd, *queryRangeStep))
 
 	case querySeriesCmd.FullCommand():
 		os.Exit(QuerySeries(*querySeriesServer, *querySeriesMatch, *querySeriesBegin, *querySeriesEnd))
@@ -413,7 +414,7 @@ func QueryInstant(url string, query string) int {
 }
 
 // QueryRange performs a range query against a Prometheus server.
-func QueryRange(url string, query string, start string, end string) int {
+func QueryRange(url, query, start, end string, step time.Duration) int {
 	config := api.Config{
 		Address: url,
 	}
@@ -450,9 +451,11 @@ func QueryRange(url string, query string, start string, end string) int {
 		fmt.Fprintln(os.Stderr, "start time is not before end time")
 	}
 
-	resolution := math.Max(math.Floor(etime.Sub(stime).Seconds()/250), 1)
-	// Convert seconds to nanoseconds such that time.Duration parses correctly.
-	step := time.Duration(resolution * 1e9)
+	if step == 0 {
+		resolution := math.Max(math.Floor(etime.Sub(stime).Seconds()/250), 1)
+		// Convert seconds to nanoseconds such that time.Duration parses correctly.
+		step = time.Duration(resolution) * time.Second
+	}
 
 	// Run query against client.
 	api := v1.NewAPI(c)
diff --git a/cmd/promtool/main_test.go b/cmd/promtool/main_test.go
new file mode 100644
index 000000000..ae71a3f63
--- /dev/null
+++ b/cmd/promtool/main_test.go
@@ -0,0 +1,75 @@
+// Copyright 2018 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+)
+
+func TestQueryRange(t *testing.T) {
+	s, getURL := mockServer(200, `{"status": "success", "data": {"resultType": "matrix", "result": []}}`)
+	defer s.Close()
+
+	exitCode := QueryRange(s.URL, "up", "0", "300", 0)
+	expectedPath := "/api/v1/query_range"
+	if getURL().Path != expectedPath {
+		t.Errorf("unexpected URL path %s (wanted %s)", getURL().Path, expectedPath)
+	}
+	actual := getURL().Query().Get("query")
+	if actual != "up" {
+		t.Errorf("unexpected value %s for query", actual)
+	}
+	actual = getURL().Query().Get("step")
+	if actual != "1.000" {
+		t.Errorf("unexpected value %s for step", actual)
+	}
+	if exitCode > 0 {
+		t.Error()
+	}
+
+	exitCode = QueryRange(s.URL, "up", "0", "300", 10*time.Millisecond)
+	if getURL().Path != expectedPath {
+		t.Errorf("unexpected URL path %s (wanted %s)", getURL().Path, expectedPath)
+	}
+	actual = getURL().Query().Get("query")
+	if actual != "up" {
+		t.Errorf("unexpected value %s for query", actual)
+	}
+	actual = getURL().Query().Get("step")
+	if actual != "0.010" {
+		t.Errorf("unexpected value %s for step", actual)
+	}
+	if exitCode > 0 {
+		t.Error()
+	}
+}
+
+func mockServer(code int, body string) (*httptest.Server, func() *url.URL) {
+	var u *url.URL
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		u = r.URL
+		w.WriteHeader(code)
+		fmt.Fprintln(w, body)
+	}))
+
+	f := func() *url.URL {
+		return u
+	}
+	return server, f
+}