1 files changed, 71 insertions, 0 deletions
diff --git a/go/internal/command/command_test.go b/go/internal/command/command_test.go
new file mode 100644
index 0000000..02fc0d0
--- /dev/null
+++ b/go/internal/command/command_test.go
@@ -0,0 +1,71 @@
+package command
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/discover"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/command/fallback"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/go/internal/testhelper"
+)
+
+func TestNew(t *testing.T) {
+	testCases := []struct {
+		desc         string
+		arguments    []string
+		config       *config.Config
+		environment  map[string]string
+		expectedType interface{}
+	}{
+		{
+			desc:      "it returns a Discover command if the feature is enabled",
+			arguments: []string{},
+			config: &config.Config{
+				GitlabUrl: "http+unix://gitlab.socket",
+				Migration: config.MigrationConfig{Enabled: true, Features: []string{"discover"}},
+			},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "",
+			},
+			expectedType: &discover.Command{},
+		},
+		{
+			desc:      "it returns a Fallback command no feature is enabled",
+			arguments: []string{},
+			config: &config.Config{
+				GitlabUrl: "http+unix://gitlab.socket",
+				Migration: config.MigrationConfig{Enabled: false},
+			},
+			environment: map[string]string{
+				"SSH_CONNECTION":       "1",
+				"SSH_ORIGINAL_COMMAND": "",
+			},
+			expectedType: &fallback.Command{},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			restoreEnv := testhelper.TempEnv(tc.environment)
+			defer restoreEnv()
+
+			command, err := New(tc.arguments, tc.config)
+
+			assert.NoError(t, err)
+			assert.IsType(t, tc.expectedType, command)
+		})
+	}
+}
+
+func TestFailingNew(t *testing.T) {
+	t.Run("It returns an error when SSH_CONNECTION is not set", func(t *testing.T) {
+		restoreEnv := testhelper.TempEnv(map[string]string{})
+		defer restoreEnv()
+
+		_, err := New([]string{}, &config.Config{})
+
+		assert.Error(t, err, "Only ssh allowed")
+	})
+}