From 2ecda75dc3d41e68e223b188fa54d4a304fce404 Mon Sep 17 00:00:00 2001
From: Denis Krivak <tetafro@gmail.com>
Date: Wed, 13 Sep 2023 09:01:04 +0200
Subject: [PATCH] Add tests.

---
 collector_test.go  |  64 ++++++++++++++++++++++++++
 connectbox_test.go | 109 ++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 167 insertions(+), 6 deletions(-)

diff --git a/collector_test.go b/collector_test.go
index a60d6b9..2840965 100644
--- a/collector_test.go
+++ b/collector_test.go
@@ -3,6 +3,8 @@ package main
 import (
 	"context"
 	"errors"
+	"io"
+	"log"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -18,6 +20,8 @@ func TestNewCollector(t *testing.T) {
 }
 
 func TestCollector_ServeHTTP(t *testing.T) {
+	log.SetOutput(io.Discard)
+
 	t.Run("success", func(t *testing.T) {
 		ctrl := gomock.NewController(t)
 
@@ -180,4 +184,64 @@ func TestCollector_ServeHTTP(t *testing.T) {
 
 		require.Equal(t, http.StatusInternalServerError, rec.Code)
 	})
+
+	t.Run("failed to get metrics", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+
+		metrics := NewMockMetricsClient(ctrl)
+
+		metrics.EXPECT().Login(gomock.Any()).Return(nil)
+		metrics.EXPECT().GetMetrics(gomock.Any(), FnCMSystemInfo, gomock.Any()).
+			Return(errors.New("fail"))
+		metrics.EXPECT().GetMetrics(gomock.Any(), FnLANUserTable, gomock.Any()).
+			Return(errors.New("fail"))
+		metrics.EXPECT().GetMetrics(gomock.Any(), FnCMState, gomock.Any()).
+			Return(errors.New("fail"))
+		metrics.EXPECT().Logout(gomock.Any()).Return(nil)
+
+		col := &Collector{
+			targets: map[string]MetricsClient{
+				"127.0.0.1": metrics,
+			},
+		}
+
+		req, err := http.NewRequest(http.MethodGet, "/probe?target=127.0.0.1", nil)
+		require.NoError(t, err)
+
+		rec := httptest.NewRecorder()
+		col.ServeHTTP(rec, req)
+
+		require.Equal(t, http.StatusOK, rec.Code)
+		require.Equal(t, "", rec.Body.String())
+	})
+
+	t.Run("failed to get metrics and to logout", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+
+		metrics := NewMockMetricsClient(ctrl)
+
+		metrics.EXPECT().Login(gomock.Any()).Return(nil)
+		metrics.EXPECT().GetMetrics(gomock.Any(), FnCMSystemInfo, gomock.Any()).
+			Return(errors.New("fail"))
+		metrics.EXPECT().GetMetrics(gomock.Any(), FnLANUserTable, gomock.Any()).
+			Return(errors.New("fail"))
+		metrics.EXPECT().GetMetrics(gomock.Any(), FnCMState, gomock.Any()).
+			Return(errors.New("fail"))
+		metrics.EXPECT().Logout(gomock.Any()).Return(errors.New("fail"))
+
+		col := &Collector{
+			targets: map[string]MetricsClient{
+				"127.0.0.1": metrics,
+			},
+		}
+
+		req, err := http.NewRequest(http.MethodGet, "/probe?target=127.0.0.1", nil)
+		require.NoError(t, err)
+
+		rec := httptest.NewRecorder()
+		col.ServeHTTP(rec, req)
+
+		require.Equal(t, http.StatusOK, rec.Code)
+		require.Equal(t, "", rec.Body.String())
+	})
 }
diff --git a/connectbox_test.go b/connectbox_test.go
index 33b4124..9bd1f26 100644
--- a/connectbox_test.go
+++ b/connectbox_test.go
@@ -27,13 +27,80 @@ func TestNewConnectBox(t *testing.T) {
 	})
 }
 
-func TestConnectBox_xmlRequest(t *testing.T) {
+func TestConnectBox_Logout(t *testing.T) {
+	ctx := context.Background()
+
+	connectbox := testConnectBox{
+		status: http.StatusOK,
+	}
+	server := httptest.NewServer(&connectbox)
+	defer server.Close()
+
+	client, err := NewConnectBox(server.URL, "bob", "qwerty")
+	require.NoError(t, err)
+	client.token = "abc"
+
+	err = client.Logout(ctx)
+	require.NoError(t, err)
+
+	want := "token=abc&fun=16"
+	require.Equal(t, want, connectbox.req)
+}
+
+func TestConnectBox_GetMetrics(t *testing.T) {
 	t.Run("valid response", func(t *testing.T) {
 		ctx := context.Background()
 
 		connectbox := testConnectBox{
 			status: http.StatusOK,
-			resp:   "hello, world",
+			resp:   `<?xml version="1.0"?><root><field>50</field></root>`,
+		}
+		server := httptest.NewServer(&connectbox)
+		defer server.Close()
+
+		client, err := NewConnectBox(server.URL, "bob", "qwerty")
+		require.NoError(t, err)
+		client.token = "abc"
+
+		var data struct {
+			Field string `xml:"field"`
+		}
+		err = client.GetMetrics(ctx, "100", &data)
+		require.NoError(t, err)
+		require.Equal(t, "token=abc&fun=100", connectbox.req)
+		require.Equal(t, "50", data.Field)
+	})
+
+	t.Run("invalid response", func(t *testing.T) {
+		ctx := context.Background()
+
+		connectbox := testConnectBox{
+			status: http.StatusOK,
+			resp:   `<?xml`,
+		}
+		server := httptest.NewServer(&connectbox)
+		defer server.Close()
+
+		client, err := NewConnectBox(server.URL, "bob", "qwerty")
+		require.NoError(t, err)
+		client.token = "abc"
+
+		var data struct {
+			Field string `xml:"field"`
+		}
+		err = client.GetMetrics(ctx, "100", &data)
+		require.ErrorContains(t, err, "unmarshal response")
+	})
+}
+
+func TestConnectBox_xmlRequest(t *testing.T) {
+	t.Run("valid response", func(t *testing.T) {
+		ctx := context.Background()
+
+		connectbox := testConnectBox{
+			cookies: map[string]string{"token": "def"},
+			status:  http.StatusOK,
+			resp:    "hello, world",
 		}
 		server := httptest.NewServer(&connectbox)
 		defer server.Close()
@@ -50,6 +117,7 @@ func TestConnectBox_xmlRequest(t *testing.T) {
 		require.Equal(t, want, connectbox.req)
 		want = "hello, world"
 		require.Equal(t, want, resp)
+		require.Equal(t, "def", client.getCookie("token"))
 	})
 
 	t.Run("invalid status code", func(t *testing.T) {
@@ -71,19 +139,48 @@ func TestConnectBox_xmlRequest(t *testing.T) {
 	})
 }
 
+func TestConnectBox_get(t *testing.T) {
+	ctx := context.Background()
+
+	connectbox := testConnectBox{
+		cookies: map[string]string{"token": "def"},
+		status:  http.StatusOK,
+		resp:    "hello, world",
+	}
+	server := httptest.NewServer(&connectbox)
+	defer server.Close()
+
+	client, err := NewConnectBox(server.URL, "bob", "qwerty")
+	require.NoError(t, err)
+	client.token = "abc"
+
+	resp, err := client.get(ctx, "/test")
+	require.NoError(t, err)
+	require.Equal(t, "hello, world", resp)
+	require.Equal(t, "def", client.getCookie("token"))
+}
+
 type testConnectBox struct {
-	path   string
-	req    string
-	status int
-	resp   string
+	// Save input data
+	path string
+	req  string
+	// Respond with provided data
+	cookies map[string]string
+	status  int
+	resp    string
 }
 
 func (t *testConnectBox) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// Save input data
 	t.path = r.URL.Path
 	body, err := io.ReadAll(r.Body)
 	if err == nil {
 		t.req = string(body)
 	}
+	// Respond with provided data
+	for name, val := range t.cookies {
+		http.SetCookie(w, &http.Cookie{Name: name, Value: val})
+	}
 	w.WriteHeader(t.status)
 	w.Write([]byte(t.resp))
 }