refactor(event_hubs): refactored to plugin/commons

plugins/common/eh/eh.go

@@ -0,0 +1,90 @@
+package eh
+
+import (
+	"context"
+	"time"
+
+	eventhub "github.com/Azure/azure-event-hubs-go/v3"
+
+	"github.com/influxdata/telegraf"
+	"github.com/influxdata/telegraf/config"
+)
+
+var sampleConfig string
+
+/*
+** Wrapper interface for eventhub.Hub
+ */
+
+type EventHubInterface interface {
+	GetHub(s string) error
+	Close(ctx context.Context) error
+	SendBatch(ctx context.Context, iterator eventhub.BatchIterator, opts ...eventhub.BatchOption) error
+}
+
+type EventHub struct {
+	hub *eventhub.Hub
+}
+
+func (eh *EventHub) GetHub(s string) error {
+	hub, err := eventhub.NewHubFromConnectionString(s)
+
+	if err != nil {
+		return err
+	}
+
+	eh.hub = hub
+
+	return nil
+}
+
+func (eh *EventHub) Close(ctx context.Context) error {
+	return eh.hub.Close(ctx)
+}
+
+func (eh *EventHub) SendBatch(ctx context.Context, iterator eventhub.BatchIterator, opts ...eventhub.BatchOption) error {
+	return eh.hub.SendBatch(ctx, iterator, opts...)
+}
+
+/* End wrapper interface */
+
+type EventHubs struct {
+	Log              telegraf.Logger `toml:"-"`
+	ConnectionString string          `toml:"connection_string"`
+	Timeout          config.Duration `toml:"timeout"`
+	PartitionKey     string          `toml:"partition_key"`
+	MaxMessageSize   int             `toml:"max_message_size"`
+
+	Hub EventHubInterface
+}
+
+func (*EventHubs) SampleConfig() string {
+	return sampleConfig
+}
+
+func (e *EventHubs) Init() error {
+	err := e.Hub.GetHub(e.ConnectionString)
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (e *EventHubs) Connect() error {
+	return nil
+}
+
+func (e *EventHubs) Close() error {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(e.Timeout))
+	defer cancel()
+
+	err := e.Hub.Close(ctx)
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

plugins/common/eh/eh_test.go

@@ -0,0 +1,34 @@
+package eh
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/influxdata/telegraf/config"
+)
+
+/*
+** Integration test (requires an Event Hubs instance)
+ */
+
+func TestInit(t *testing.T) {
+	if testing.Short() {
+		t.Skip("Skipping integration test in short mode")
+	}
+
+	connstring := "Endpoint=sb://telegraf.servicebus.windows.net/;SharedAccessKeyName=RootManageSharedAccessKey;SharedAccessKey=+"
+
+	// Configure the plugin to target the newly created hub
+	e := &EventHubs{
+		Hub:              &EventHub{},
+		ConnectionString: connstring,
+		Timeout:          config.Duration(time.Second * 5),
+	}
+
+	err := e.Init()
+	// Verify that we can connect to Event Hubs
+	require.NoError(t, err)
+	e.Close()
+}

plugins/outputs/event_hubs/event_hubs.go

@@ -7,100 +7,47 @@ import (
 	"time"
 
 	eventhub "github.com/Azure/azure-event-hubs-go/v3"
-
 	"github.com/influxdata/telegraf"
 	"github.com/influxdata/telegraf/config"
+	eh_commons "github.com/influxdata/telegraf/plugins/common/eh"
+
 	"github.com/influxdata/telegraf/plugins/outputs"
 )
 
-//go:embed sample.conf
-var sampleConfig string
-
-/*
-** Wrapper interface for eventhub.Hub
- */
-
-type EventHubInterface interface {
-	GetHub(s string) error
-	Close(ctx context.Context) error
-	SendBatch(ctx context.Context, iterator eventhub.BatchIterator, opts ...eventhub.BatchOption) error
-}
-
-type eventHub struct {
-	hub *eventhub.Hub
-}
-
-func (eh *eventHub) GetHub(s string) error {
-	hub, err := eventhub.NewHubFromConnectionString(s)
-
-	if err != nil {
-		return err
-	}
-
-	eh.hub = hub
-
-	return nil
-}
-
-func (eh *eventHub) Close(ctx context.Context) error {
-	return eh.hub.Close(ctx)
-}
-
-func (eh *eventHub) SendBatch(ctx context.Context, iterator eventhub.BatchIterator, opts ...eventhub.BatchOption) error {
-	return eh.hub.SendBatch(ctx, iterator, opts...)
-}
-
 /* End wrapper interface */
 
 type EventHubs struct {
-	Log              telegraf.Logger `toml:"-"`
-	ConnectionString string          `toml:"connection_string"`
-	Timeout          config.Duration `toml:"timeout"`
-	PartitionKey     string          `toml:"partition_key"`
-	MaxMessageSize   int             `toml:"max_message_size"`
-
-	Hub          EventHubInterface
-	batchOptions []eventhub.BatchOption
+	eh_commons.EventHubs
 	serializer   telegraf.Serializer
+	batchOptions []eventhub.BatchOption
 }
 
 const (
 	defaultRequestTimeout = time.Second * 30
 )
 
-func (*EventHubs) SampleConfig() string {
+//go:embed sample.conf
+var sampleConfig string
+
+func (e *EventHubs) SampleConfig() string {
 	return sampleConfig
 }
 
 func (e *EventHubs) Init() error {
-	err := e.Hub.GetHub(e.ConnectionString)
-
-	if err != nil {
-		return err
-	}
 
 	if e.MaxMessageSize > 0 {
 		e.batchOptions = append(e.batchOptions, eventhub.BatchWithMaxSizeInBytes(e.MaxMessageSize))
 	}
 
-	return nil
+	return e.EventHubs.Init()
 }
 
-func (*EventHubs) Connect() error {
+func (e *EventHubs) Connect() error {
 	return nil
 }
 
 func (e *EventHubs) Close() error {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(e.Timeout))
-	defer cancel()
-
-	err := e.Hub.Close(ctx)
-
-	if err != nil {
-		return err
-	}
-
-	return nil
+	return e.EventHubs.Close()
 }
 
 func (e *EventHubs) SetSerializer(serializer telegraf.Serializer) {
@@ -146,8 +93,10 @@ func (e *EventHubs) Write(metrics []telegraf.Metric) error {
 func init() {
 	outputs.Add("event_hubs", func() telegraf.Output {
 		return &EventHubs{
-			Hub:     &eventHub{},
-			Timeout: config.Duration(defaultRequestTimeout),
+			EventHubs: eh_commons.EventHubs{
+				Hub:     &eh_commons.EventHub{},
+				Timeout: config.Duration(defaultRequestTimeout),
+			},
 		}
 	})
 }

plugins/outputs/event_hubs/event_hubs_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/influxdata/telegraf/config"
+	eh_commons "github.com/influxdata/telegraf/plugins/common/eh"
 	"github.com/influxdata/telegraf/plugins/serializers/json"
 	"github.com/influxdata/telegraf/testutil"
 )
@@ -47,14 +48,16 @@ func TestInitAndWrite(t *testing.T) {
 	require.NoError(t, serializer.Init())
 
 	mockHub := &mockEventHub{}
-	e := &EventHubs{
-		Hub:              mockHub,
-		ConnectionString: "mock",
-		Timeout:          config.Duration(time.Second * 5),
-		MaxMessageSize:   1000000,
-		serializer:       serializer,
-	}
 
+	e := EventHubs{
+		EventHubs: eh_commons.EventHubs{
+			Hub:              mockHub,
+			ConnectionString: "mock",
+			Timeout:          config.Duration(time.Second * 5),
+			MaxMessageSize:   1000000,
+		},
+		serializer: serializer,
+	}
 	mockHub.On("GetHub", mock.Anything).Return(nil).Once()
 	require.NoError(t, e.Init())
 	mockHub.AssertExpectations(t)
@@ -104,11 +107,14 @@ func TestInitAndWriteIntegration(t *testing.T) {
 	// Configure the plugin to target the newly created hub
 	serializer := &json.Serializer{}
 	require.NoError(t, serializer.Init())
-	e := &EventHubs{
-		Hub:              &eventHub{},
-		ConnectionString: testHubCS,
-		Timeout:          config.Duration(time.Second * 5),
-		serializer:       serializer,
+	e := EventHubs{
+
+		EventHubs: eh_commons.EventHubs{
+			Hub:              &eh_commons.EventHub{},
+			ConnectionString: testHubCS,
+			Timeout:          config.Duration(time.Second * 5),
+		},
+		serializer: serializer,
 	}
 
 	// Verify that we can connect to Event Hubs