From 582bba5f94fe113051c318ba8c6758cef5b10372 Mon Sep 17 00:00:00 2001
From: Conor Broderick <conor.broderick@robustperception.io>
Date: Wed, 12 Jul 2017 17:09:59 +0100
Subject: [PATCH] Added test for tcp timeout (#192)

---
 http_test.go | 34 ++++++++++++++++++++++++++++++++++
 main_test.go | 42 ------------------------------------------
 tcp_test.go  | 31 +++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 42 deletions(-)
 delete mode 100644 main_test.go

diff --git a/http_test.go b/http_test.go
index eec5c1b..5e0a632 100644
--- 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
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 7c2367c..0551792 100644
--- 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
+}
-- 
2.25.1