From f572d5bcf3283de99e7e83c31c56c221c1034f40 Mon Sep 17 00:00:00 2001
From: colmsnowplow <colm@snowplowanalytics.com>
Date: Wed, 21 Jul 2021 16:55:23 +0100
Subject: [PATCH] Add more thorough check on auth vars, and tests for the check

---
 pkg/target/eventhub.go      | 19 +++++++---
 pkg/target/eventhub_test.go | 70 +++++++++++++++++++++++++++++++++++++
 2 files changed, 85 insertions(+), 4 deletions(-)
 create mode 100644 pkg/target/eventhub_test.go

diff --git a/pkg/target/eventhub.go b/pkg/target/eventhub.go
index 1242fbf3..65f66333 100644
--- a/pkg/target/eventhub.go
+++ b/pkg/target/eventhub.go
@@ -51,10 +51,21 @@ type EventHubTarget struct {
 // NewEventHubTarget creates a new client for writing messages to Azure EventHub
 func NewEventHubTarget(cfg *EventHubConfig) (*EventHubTarget, error) {
 
-	keyName, keyNamePresent := os.LookupEnv("EVENTHUB_KEY_NAME")
-	keyValue, keyValuePresent := os.LookupEnv("EVENTHUB_KEY_VALUE")
-	if !(keyNamePresent && keyValuePresent) {
-		return nil, errors.Errorf("Error initialising EventHub client: Env vars %s and %s are required for authentication.", keyName, keyValue)
+	_, keyNamePresent := os.LookupEnv("EVENTHUB_KEY_NAME")
+	_, keyValuePresent := os.LookupEnv("EVENTHUB_KEY_VALUE")
+
+	_, connStringPresent := os.LookupEnv("EVENTHUB_CONNECTION_STRING")
+
+	_, tenantIDPresent := os.LookupEnv("AZURE_TENANT_ID")
+	_, clientIDPresent := os.LookupEnv("AZURE_CLIENT_ID")
+
+	_, clientSecretPresent := os.LookupEnv("AZURE_CLIENT_SECRET")
+
+	_, azCertPathPresent := os.LookupEnv("AZURE_CERTIFICATE_PATH")
+	_, azCertPwrdPresent := os.LookupEnv("AZURE_CERTIFICATE_PASSWORD")
+
+	if !(connStringPresent || (keyNamePresent && keyValuePresent) || (tenantIDPresent && clientIDPresent && ((azCertPathPresent && azCertPwrdPresent) || clientSecretPresent))) {
+		return nil, errors.Errorf("Error initialising EventHub client: No valid combination of authentication Env vars found. https://pkg.go.dev/github.com/Azure/azure-event-hubs-go#NewHubWithNamespaceNameAndEnvironment")
 	}
 
 	hub, err := eventhub.NewHubWithNamespaceNameAndEnvironment(cfg.EventHubNamespace, cfg.EventHubName)
diff --git a/pkg/target/eventhub_test.go b/pkg/target/eventhub_test.go
new file mode 100644
index 00000000..9f9efb50
--- /dev/null
+++ b/pkg/target/eventhub_test.go
@@ -0,0 +1,70 @@
+package target
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var cfg = EventHubConfig{
+	EventHubNamespace: "test",
+	EventHubName:      "test",
+}
+
+func unsetEverything() {
+	os.Unsetenv("EVENTHUB_KEY_NAME")
+	os.Unsetenv("EVENTHUB_KEY_VALUE")
+
+	os.Unsetenv("EVENTHUB_CONNECTION_STRING")
+
+	os.Unsetenv("AZURE_TENANT_ID")
+	os.Unsetenv("AZURE_CLIENT_ID")
+
+	os.Unsetenv("AZURE_CLIENT_SECRET")
+
+	os.Unsetenv("AZURE_CERTIFICATE_PATH")
+	os.Unsetenv("AZURE_CERTIFICATE_PASSWORD")
+}
+
+func TestNewEventHubTarget_KeyValue(t *testing.T) {
+	assert := assert.New(t)
+
+	unsetEverything()
+
+	// Test that we can initialise a client with Key and Value
+	defer os.Unsetenv("EVENTHUB_KEY_NAME")
+	defer os.Unsetenv("EVENTHUB_KEY_VALUE")
+
+	os.Setenv("EVENTHUB_KEY_NAME", "fake")
+	os.Setenv("EVENTHUB_KEY_VALUE", "fake")
+
+	tgt, err := NewEventHubTarget(&cfg)
+	assert.Nil(err)
+	assert.NotNil(tgt)
+}
+
+func TestNewEventHubTarget_ConnString(t *testing.T) {
+	assert := assert.New(t)
+
+	unsetEverything()
+
+	// Test that we can initialise a client with Connection String
+	defer os.Unsetenv("EVENTHUB_CONNECTION_STRING")
+
+	os.Setenv("EVENTHUB_CONNECTION_STRING", "Endpoint=sb://test.servicebus.windows.net/;SharedAccessKeyName=fake;SharedAccessKey=fake")
+
+	tgt, err := NewEventHubTarget(&cfg)
+	assert.Nil(err)
+	assert.NotNil(tgt)
+}
+
+func TestNewEventHubTarget_Failure(t *testing.T) {
+	assert := assert.New(t)
+
+	unsetEverything()
+
+	tgt, err := NewEventHubTarget(&cfg)
+	assert.Equal("Error initialising EventHub client: No valid combination of authentication Env vars found. https://pkg.go.dev/github.com/Azure/azure-event-hubs-go#NewHubWithNamespaceNameAndEnvironment", err.Error())
+	assert.Nil(tgt)
+}