Fix DCat color test.

10 files changed, 37 insertions, 36 deletions

diff --git a/internal/config/client.go b/internal/config/client.go
index 86f97f0..9f4df97 100644
--- a/internal/config/client.go
+++ b/internal/config/client.go
@@ -104,9 +104,6 @@ type termColors struct {
 type ClientConfig struct {
 	TermColorsEnable bool       `json:",omitempty"`
 	TermColors       termColors `json:",omitempty"`
-	// When unit testing in Jenkins you don't want to touch files in ~jenkins
-	// during integration tests really.
-	SSHDontAddHostsToKnownHostsFile bool `json:",omitempty"`
 }
 
 // Create a new default client configuration.
diff --git a/internal/config/env.go b/internal/config/env.go
index 88b831d..804a10a 100644
--- a/internal/config/env.go
+++ b/internal/config/env.go
@@ -2,6 +2,26 @@ package config
 
 import "os"
 
+// Env returns true when a given environment variable is set to "yes".
 func Env(env string) bool {
 	return "yes" == os.Getenv(env)
 }
+
+// Hostname returns the current hostname. It can be overriden with
+// DTAIL_HOSTNAME_OVERRIDE environment variable (useful for integration tests).
+func Hostname() (string, error) {
+	hostname := os.Getenv("DTAIL_HOSTNAME_OVERRIDE")
+	if len(hostname) > 0 {
+		return hostname, nil
+	}
+	return os.Hostname()
+}
+
+// SSHKnownHostsFile returns the known hosts file path (useful for integration tests)
+func SSHKnownHostsFile() string {
+	if len(os.Getenv("DTAIL_SSH_KNOWN_HOSTS_FILE")) > 0 {
+		return os.Getenv("DTAIL_SSH_KNOWN_HOSTS_FILE")
+	} else {
+		return os.Getenv("HOME") + "/.ssh/known_hosts"
+	}
+}
diff --git a/internal/config/initializer.go b/internal/config/initializer.go
index 4d6a73b..8d617f2 100644
--- a/internal/config/initializer.go
+++ b/internal/config/initializer.go
@@ -82,9 +82,9 @@ func (in *initializer) transformConfig(sourceProcess source.Source, args *Args,
 
 // There are some special options which can be set by environment variable.
 func (in *initializer) readEnvironmentVars() {
-	if Env("DTAIL_SSH_DONT_ADD_HOSTS_TO_KNOWNHOSTS_FILE") ||
-		Env("DTAIL_JENKINS") {
-		in.Client.SSHDontAddHostsToKnownHostsFile = true
+	if Env("DTAIL_RUN_INTEGRATION_TESTS") {
+		os.Setenv("DTAIL_HOSTNAME_OVERRIDE", "integrationtest")
+		os.Setenv("DTAIL_SSH_KNOWN_HOSTS_FILE", "./known_hosts")
 	}
 }
 
diff --git a/internal/io/dlog/dlog.go b/internal/io/dlog/dlog.go
index 5e0c3a1..ff2cef4 100644
--- a/internal/io/dlog/dlog.go
+++ b/internal/io/dlog/dlog.go
@@ -78,7 +78,7 @@ type DLog struct {
 
 // new creates a new DTail logger.
 func new(sourceProcess, sourcePackage source.Source) *DLog {
-	hostname, err := os.Hostname()
+	hostname, err := config.Hostname()
 	if err != nil {
 		panic(err)
 	}
diff --git a/internal/mapr/logformat/parser.go b/internal/mapr/logformat/parser.go
index 129081d..d6aac8c 100644
--- a/internal/mapr/logformat/parser.go
+++ b/internal/mapr/logformat/parser.go
@@ -3,11 +3,11 @@ package logformat
 import (
 	"errors"
 	"fmt"
-	"os"
 	"reflect"
 	"strings"
 	"time"
 
+	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/mapr"
 )
 
@@ -26,7 +26,7 @@ type Parser struct {
 
 // NewParser returns a new log parser.
 func NewParser(logFormatName string, query *mapr.Query) (*Parser, error) {
-	hostname, err := os.Hostname()
+	hostname, err := config.Hostname()
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/mapr/server/aggregate.go b/internal/mapr/server/aggregate.go
index 11c9ee5..4162828 100644
--- a/internal/mapr/server/aggregate.go
+++ b/internal/mapr/server/aggregate.go
@@ -2,7 +2,6 @@ package server
 
 import (
 	"context"
-	"os"
 	"strings"
 	"time"
 
@@ -38,7 +37,7 @@ func NewAggregate(queryStr string) (*Aggregate, error) {
 		return nil, err
 	}
 
-	fqdn, err := os.Hostname()
+	fqdn, err := config.Hostname()
 	if err != nil {
 		dlog.Common.Error(err)
 	}
diff --git a/internal/server/handlers/healthhandler.go b/internal/server/handlers/healthhandler.go
index 6dd9872..e7f7762 100644
--- a/internal/server/handlers/healthhandler.go
+++ b/internal/server/handlers/healthhandler.go
@@ -2,10 +2,10 @@ package handlers
 
 import (
 	"context"
-	"os"
 	"strings"
 
 	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/io/line"
 	"github.com/mimecast/dtail/internal/lcontext"
@@ -32,7 +32,7 @@ func NewHealthHandler(user *user.User) *HealthHandler {
 	}
 	h.handleCommandCb = h.handleHealthCommand
 
-	fqdn, err := os.Hostname()
+	fqdn, err := config.Hostname()
 	if err != nil {
 		dlog.Server.FatalPanic(err)
 	}
diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index 36574a9..75a8acc 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -2,10 +2,10 @@ package handlers
 
 import (
 	"context"
-	"os"
 	"strings"
 
 	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/io/line"
 	"github.com/mimecast/dtail/internal/lcontext"
@@ -43,7 +43,7 @@ func NewServerHandler(user *user.User, catLimiter,
 	}
 	h.handleCommandCb = h.handleUserCommand
 
-	fqdn, err := os.Hostname()
+	fqdn, err := config.Hostname()
 	if err != nil {
 		dlog.Server.FatalPanic(err)
 	}
diff --git a/internal/ssh/client/authmethods.go b/internal/ssh/client/authmethods.go
index 37f8382..2ee32ad 100644
--- a/internal/ssh/client/authmethods.go
+++ b/internal/ssh/client/authmethods.go
@@ -29,19 +29,14 @@ func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
 	privateKeyPath string) ([]gossh.AuthMethod, HostKeyCallback) {
 
 	var sshAuthMethods []gossh.AuthMethod
-	knownHostsPath := os.Getenv("HOME") + "/.ssh/known_hosts"
-	knownHostsCallback, err := NewKnownHostsCallback(knownHostsPath, trustAllHosts,
+	knownHostsFile := config.SSHKnownHostsFile()
+	knownHostsCallback, err := NewKnownHostsCallback(knownHostsFile, trustAllHosts,
 		throttleCh)
 	if err != nil {
-		dlog.Client.FatalPanic(knownHostsPath, err)
+		dlog.Client.FatalPanic(knownHostsFile, err)
 	}
-	dlog.Client.Debug("initKnownHostsAuthMethods", "Added known hosts file path", knownHostsPath)
-	/*
-		if config.Client.ExperimentalFeaturesEnable {
-			sshAuthMethods = append(sshAuthMethods, gossh.Password("experimental feature test"))
-			dlog.Client.Debug("initKnownHostsAuthMethods", "Added experimental method to list of auth methods")
-		}
-	*/
+
+	dlog.Client.Debug("initKnownHostsAuthMethods", "Added known hosts file path", knownHostsFile)
 
 	// First try to read custom private key path.
 	if privateKeyPath != "" {
@@ -100,11 +95,7 @@ func initKnownHostsAuthMethods(trustAllHosts bool, throttleCh chan struct{},
 	dlog.Client.Debug("initKnownHostsAuthMethods", "Unable to use private key",
 		privateKeyPath, err)
 
-	// This is only a panic when we expect to do something about it.
-	if !config.Client.SSHDontAddHostsToKnownHostsFile {
-		dlog.Client.FatalPanic("Unable to find private SSH key information")
-	}
-
+	dlog.Client.FatalPanic("Unable to find private SSH key information")
 	// Never reach this point.
 	return sshAuthMethods, knownHostsCallback
 }
diff --git a/internal/ssh/client/knownhostscallback.go b/internal/ssh/client/knownhostscallback.go
index 2aa0168..dd58925 100644
--- a/internal/ssh/client/knownhostscallback.go
+++ b/internal/ssh/client/knownhostscallback.go
@@ -10,7 +10,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/io/prompt"
 
@@ -216,11 +215,6 @@ func (c KnownHostsCallback) promptAddHosts(hosts []unknownHost) {
 func (c KnownHostsCallback) trustHosts(hosts []unknownHost) {
 	tmpKnownHostsPath := fmt.Sprintf("%s.tmp", c.knownHostsPath)
 
-	if config.Client.SSHDontAddHostsToKnownHostsFile {
-		dlog.Common.Verbose("Not adding hosts to known hosts file, as disabled by config")
-		return
-	}
-
 	newFd, err := os.OpenFile(tmpKnownHostsPath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600)
 	if err != nil {
 		panic(fmt.Sprintf("%s: %s", tmpKnownHostsPath, err.Error()))