From e31e833eace4d9b88c69a7430108cec9f27e96f5 Mon Sep 17 00:00:00 2001
From: Suraj Shirvankar <surajshirvankar@gmail.com>
Date: Thu, 8 Aug 2024 18:09:12 +0200
Subject: [PATCH] refactor: add tests for bootstrap redirect

Signed-off-by: Suraj Shirvankar <surajshirvankar@gmail.com>
---
 sztp-agent/pkg/secureagent/daemon.go      |  7 +++
 sztp-agent/pkg/secureagent/daemon_test.go | 77 +++++++++++++++++++++++
 2 files changed, 84 insertions(+)

diff --git a/sztp-agent/pkg/secureagent/daemon.go b/sztp-agent/pkg/secureagent/daemon.go
index cc0fa44..2064d90 100644
--- a/sztp-agent/pkg/secureagent/daemon.go
+++ b/sztp-agent/pkg/secureagent/daemon.go
@@ -13,6 +13,7 @@ import (
 	"encoding/asn1"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"log"
 	"net/url"
@@ -114,6 +115,12 @@ func (a *Agent) doHandleBootstrapRedirect() error {
 	addr := a.BootstrapServerRedirectInfo.IetfSztpConveyedInfoRedirectInformation.BootstrapServer[0].Address
 	port := a.BootstrapServerRedirectInfo.IetfSztpConveyedInfoRedirectInformation.BootstrapServer[0].Port
 
+	if addr == "" {
+		return errors.New("invalid redirect address")
+	}
+	if port <= 0 {
+		return errors.New("invalid port")
+	}
 	// Change URL to point to new redirect IP and PORT
 	u, err := url.Parse(a.GetBootstrapURL())
 	if err != nil {
diff --git a/sztp-agent/pkg/secureagent/daemon_test.go b/sztp-agent/pkg/secureagent/daemon_test.go
index 17578f3..4e9ba1e 100644
--- a/sztp-agent/pkg/secureagent/daemon_test.go
+++ b/sztp-agent/pkg/secureagent/daemon_test.go
@@ -164,6 +164,83 @@ func deleteTempTestFile(file string) {
 	}
 }
 
+func TestAgent_doHandleBootstrapRedirect(t *testing.T) {
+	type fields struct {
+		InputBootstrapURL           string
+		BootstrapServerRedirectInfo BootstrapServerRedirectInfo
+	}
+	tests := []struct {
+		name                 string
+		fields               fields
+		wantErr              bool
+		expectedBootstrapURL string
+	}{
+		{
+			name: "Fail test with invalid address",
+			fields: fields{
+				InputBootstrapURL: "",
+				BootstrapServerRedirectInfo: BootstrapServerRedirectInfo{
+					IetfSztpConveyedInfoRedirectInformation: struct {
+						BootstrapServer []struct {
+							Address     string `json:"address"`
+							Port        int    `json:"port"`
+							TrustAnchor string `json:"trust-anchor"`
+						} `json:"bootstrap-server"`
+					}{
+						BootstrapServer: []struct {
+							Address     string `json:"address"`
+							Port        int    `json:"port"`
+							TrustAnchor string `json:"trust-anchor"`
+						}{{
+							Address: "",
+							Port:    0,
+						}},
+					},
+				},
+			},
+			wantErr:              true,
+			expectedBootstrapURL: "",
+		},
+		{
+			name: "Fail test with invalid port",
+			fields: fields{
+				InputBootstrapURL: "",
+				BootstrapServerRedirectInfo: BootstrapServerRedirectInfo{
+					IetfSztpConveyedInfoRedirectInformation: struct {
+						BootstrapServer []struct {
+							Address     string `json:"address"`
+							Port        int    `json:"port"`
+							TrustAnchor string `json:"trust-anchor"`
+						} `json:"bootstrap-server"`
+					}{
+						BootstrapServer: []struct {
+							Address     string `json:"address"`
+							Port        int    `json:"port"`
+							TrustAnchor string `json:"trust-anchor"`
+						}{{
+							Address: "8.8.8.8",
+							Port:    -1000,
+						}},
+					},
+				},
+			},
+			wantErr:              true,
+			expectedBootstrapURL: "",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			a := &Agent{
+				BootstrapURL:                tt.fields.InputBootstrapURL,
+				BootstrapServerRedirectInfo: tt.fields.BootstrapServerRedirectInfo,
+			}
+			if err := a.doHandleBootstrapRedirect(); (err != nil) != tt.wantErr {
+				t.Errorf("doHandleBootstrapRedirect() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}
+
 //nolint:funlen
 func TestAgent_doReqBootstrap(t *testing.T) {
 	var output []byte