refactor: add more test for ssh.go

Signed-off-by: Suraj Shirvankar <[email protected]>
opiproject · Aug 16, 2024 · 8a96a2a · 8a96a2a
1 parent aa72379
commit 8a96a2a
Showing 1 changed file with 17 additions and 1 deletion.
diff --git a/sztp-agent/pkg/secureagent/ssh_test.go b/sztp-agent/pkg/secureagent/ssh_test.go
@@ -12,6 +12,7 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 		file      string
 		content   string
 		Algorithm string
+		KeyOnly   bool
 	}
 	tests := []struct {
 		name string
@@ -24,6 +25,7 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 				file:      "/tmp/test.pub",
 				content:   "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR",
 				Algorithm: "ssh-ed25519",
+				KeyOnly:   false,
 			},
 			want: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR",
 		},
@@ -33,6 +35,7 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 				file:      "/tmp/test.pub",
 				content:   "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR comment",
 				Algorithm: "ssh-ed25519",
+				KeyOnly:   false,
 			},
 			want: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR",
 		},
@@ -42,6 +45,7 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 				file:      "/tmp/test.pub",
 				content:   "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR comment error",
 				Algorithm: "ssh-ed25519",
+				KeyOnly:   false,
 			},
 			want: "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR",
 		},
@@ -50,6 +54,7 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 			args: args{
 				file:    "/tmp/test.pub",
 				content: "ssh-ed25519",
+				KeyOnly: false,
 			},
 			want: "ssh-ed25519",
 		},
@@ -58,17 +63,28 @@ func Test_readSSHHostKeyPublicFiles(t *testing.T) {
 			args: args{
 				file:    "/tmp/test.pub",
 				content: "",
+				KeyOnly: false,
 			},
 			want: "",
 		},
+		{
+			name: "Test file only key",
+			args: args{
+				file:      "/tmp/test.pub",
+				content:   "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR",
+				Algorithm: "ssh-ed25519",
+				KeyOnly:   true,
+			},
+			want: "AAAAC3NzaC1lZDI1NTE5AAAAID0mjQXlOvkM2HO5vTrSOdHOl3BGOqDiHrx8yYdbP8xR",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			if tt.args.content != "" {
 				createTempTestFile(tt.args.file, tt.args.content, true)
 			}
 			for _, key := range readSSHHostKeyPublicFiles(tt.args.file) {
-				if got := getSSHHostKeyString(key, true); !reflect.DeepEqual(got, tt.want) {
+				if got := getSSHHostKeyString(key, !tt.args.KeyOnly); !reflect.DeepEqual(got, tt.want) {
 					t.Errorf("readSSHHostKeyPublicFiles() - got: %v, want %v", got, tt.want)
 				}
 			}