Use mock MTA-STS check

EFForg · Mar 6, 2019 · 78a45a9 · 78a45a9
1 parent 43b7061
commit 78a45a9
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/checker/checker.go b/checker/checker.go
@@ -24,7 +24,7 @@ type Checker struct {
 	checkHostnameOverride func(string, string) HostnameResult
 
 	// checkMTASTSOverride is used to mock MTA-STS checks.
-	checkMTASTSOverride func(string, map[string]HostnameResult) Result
+	checkMTASTSOverride func(string, map[string]HostnameResult) *MTASTSResult
 }
 
 func (c *Checker) timeout() time.Duration {

diff --git a/checker/domain_test.go b/checker/domain_test.go
@@ -42,8 +42,10 @@ var hostnameResults = map[string]Result{
 	},
 }
 
-func mockCheckMTASTS(domain string, hostnameResults map[string]HostnameResult) Result {
-	return Result{Status: Success}
+func mockCheckMTASTS(domain string, hostnameResults map[string]HostnameResult) *MTASTSResult {
+	r := MakeMTASTSResult()
+	r.Mode = "testing"
+	return r
 }
 
 func mockLookupMX(domain string) ([]*net.MX, error) {

diff --git a/checker/mta_sts.go b/checker/mta_sts.go
@@ -177,6 +177,10 @@ func validateMTASTSMXs(policyFileMXs []string, dnsMXs map[string]HostnameResult,
 }
 
 func (c Checker) checkMTASTS(domain string, hostnameResults map[string]HostnameResult) *MTASTSResult {
+	if c.checkMTASTSOverride != nil {
+		// Allow the Checker to mock this function.
+		return c.checkMTASTSOverride(domain, hostnameResults)
+	}
 	result := MakeMTASTSResult()
 	result.addCheck(checkMTASTSRecord(domain))
 	policyResult, policy, policyMap := checkMTASTSPolicyFile(domain, hostnameResults)