refactor: add tests for bootstrap redirect

Signed-off-by: Suraj Shirvankar <[email protected]>
opiproject · Aug 16, 2024 · dae834d · dae834d
1 parent 8a96a2a
commit dae834d
Show file tree

Hide file tree

Showing 2 changed files with 84 additions and 0 deletions.
diff --git 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
@@ -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