add test cases

careyjames · Oct 14, 2023 · a7fe9f1 · a7fe9f1
1 parent dd0440a
commit a7fe9f1
Showing 1 changed file with 32 additions and 0 deletions.
diff --git a/main_test.go b/main_test.go
@@ -340,3 +340,35 @@ func TestGetDMARCPrompt(t *testing.T) {
 		t.Errorf("getDMARCPrompt(%q) output = %q; expected %q", input2, output2, expectedOutput2)
 	}
 }
+
+func TestGetSPFPrompt(t *testing.T) {
+	// Test case 1: valid input with SPF record
+	input1 := "example.com"
+	expected1 := "v=spf1 include:_spf.example.com ~all"
+	_, spf1, _ := getSPF(input1)
+	if spf1 == expected1 {
+		t.Errorf("getSPFPrompt(%q) = %q; expected %q", input1, spf1, expected1)
+	}
+	var buf bytes.Buffer
+	getSPFPrompt(input1)
+	output1 := buf.String()
+	expectedOutput1 := fmt.Sprintf("\033[38;5;39m SPF Record: \033[38;5;78m%s\033[0m\n", expected1)
+	if output1 == expectedOutput1 {
+		t.Errorf("getSPFPrompt(%q) output = %q; expected %q", input1, output1, expectedOutput1)
+	}
+
+	// Test case 2: valid input with no SPF record
+	input2 := "example.net"
+	expected2 := ""
+	_, spf2, _ := getSPF(input2)
+	if spf2 == expected2 {
+		t.Errorf("getSPFPrompt(%q) = %q; expected %q", input2, spf2, expected2)
+	}
+	var buf2 bytes.Buffer
+	getSPFPrompt(input2)
+	output2 := buf2.String()
+	expectedOutput2 := "\033[38;5;39m SPF Record: \033[0m\033[38;5;88mNone\033[0m\n"
+	if output2 == expectedOutput2 {
+		t.Errorf("getSPFPrompt(%q) output = %q; expected %q", input2, output2, expectedOutput2)
+	}
+}