Added test for tcp timeout (#192)

diff --git a/http_test.go b/http_test.go
index eec5c1bbf91b91229459326fff6236473d2a8494..5e0a6325f1c9bc84602f7db7bac6aaaa158b280c 100644 (file)
--- a/http_test.go
+++ b/http_test.go
@@ -483,3 +483,37 @@ func TestTLSConfigIsIgnoredForPlainHTTP(t *testing.T) {
        }
        checkRegistryResults(expectedResults, mfs, t)
 }
+
+var c = &Config{
+       Modules: map[string]Module{
+               "http_2xx": Module{
+                       Prober:  "http",
+                       Timeout: 10 * time.Second,
+               },
+       },
+}
+
+func TestPrometheusTimeoutHTTP(t *testing.T) {
+
+       ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+               time.Sleep(2 * time.Second)
+       }))
+       defer ts.Close()
+
+       req, err := http.NewRequest("GET", "?target="+ts.URL, nil)
+       if err != nil {
+               t.Fatal(err)
+       }
+       req.Header.Set("X-Prometheus-Scrape-Timeout-Seconds", "1")
+
+       rr := httptest.NewRecorder()
+       handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+               probeHandler(w, r, c)
+       })
+
+       handler.ServeHTTP(rr, req)
+
+       if status := rr.Code; status != http.StatusOK {
+               t.Errorf("probe request handler returned wrong status code: %v, want %v", status, http.StatusOK)
+       }
+}

diff --git a/main_test.go b/main_test.go
deleted file mode 100644 (file)
index 740cb81..0000000
--- a/main_test.go
+++ /dev/null
@@ -1,42 +0,0 @@
-package main
-
-import (
-       "net/http"
-       "net/http/httptest"
-       "testing"
-       "time"
-)
-
-var c = &Config{
-       Modules: map[string]Module{
-               "http_2xx": Module{
-                       Prober:  "http",
-                       Timeout: 10 * time.Second,
-               },
-       },
-}
-
-func TestPrometheusTimeout(t *testing.T) {
-
-       ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-               time.Sleep(2 * time.Second)
-       }))
-       defer ts.Close()
-
-       req, err := http.NewRequest("GET", "?target="+ts.URL, nil)
-       if err != nil {
-               t.Fatal(err)
-       }
-       req.Header.Set("X-Prometheus-Scrape-Timeout-Seconds", "1")
-
-       rr := httptest.NewRecorder()
-       handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-               probeHandler(w, r, c)
-       })
-
-       handler.ServeHTTP(rr, req)
-
-       if status := rr.Code; status != http.StatusOK {
-               t.Errorf("probe request handler returned wrong status code: %v, want %v", status, http.StatusOK)
-       }
-}

diff --git a/tcp_test.go b/tcp_test.go
index 7c2367c508114485e5afbfb68c239d8c329e4535..055179209d965e45409ac73cd45f311777621adb 100644 (file)
--- a/tcp_test.go
+++ b/tcp_test.go
@@ -330,3 +330,34 @@ func TestTCPConnectionProtocol(t *testing.T) {
        }
        checkRegistryResults(expectedResults, mfs, t)
 }
+
+func TestPrometheusTimeoutTCP(t *testing.T) {
+       ln, err := net.Listen("tcp", "localhost:0")
+       if err != nil {
+               t.Fatalf("Error listening on socket: %s", err)
+       }
+       defer ln.Close()
+
+       ch := make(chan (struct{}))
+       go func() {
+               conn, err := ln.Accept()
+               if err != nil {
+                       t.Fatalf("Error accepting on socket: %s", err)
+               }
+               conn.Close()
+               ch <- struct{}{}
+       }()
+       testCTX, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+       defer cancel()
+       registry := prometheus.NewRegistry()
+       if probeTCP(testCTX, ln.Addr().String(), Module{TCP: TCPProbe{
+               QueryResponse: []QueryResponse{
+                       {
+                               Expect: "SSH-2.0-(OpenSSH_6.9p1) Debian-2",
+                       },
+               },
+       }}, registry) {
+               t.Fatalf("TCP module succeeded, expected timeout failure.")
+       }
+       <-ch
+}