From 453aacb90f3f7b4e1762887a246e06f90cd275ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lovro=20Ma=C5=BEgon?= <lovro.mazgon@gmail.com>
Date: Wed, 8 Nov 2023 13:50:35 +0100
Subject: [PATCH] organize imports, simplify teardown

---
 source.go      | 11 ++++++-----
 source_test.go |  6 ++----
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/source.go b/source.go
index 905ff7cd..531bf90f 100644
--- a/source.go
+++ b/source.go
@@ -326,7 +326,6 @@ func (a *sourcePluginAdapter) Stop(ctx context.Context, _ cpluginv1.SourceStopRe
 }
 
 func (a *sourcePluginAdapter) Teardown(ctx context.Context, _ cpluginv1.SourceTeardownRequest) (cpluginv1.SourceTeardownResponse, error) {
-	var waitErr error // store waitErr
 	err := a.state.DoWithLock(ctx, internal.DoWithLockOptions{
 		ExpectedStates: nil, // Teardown can be called from any state
 		StateBefore:    internal.StateTearingDown,
@@ -344,6 +343,7 @@ func (a *sourcePluginAdapter) Teardown(ctx context.Context, _ cpluginv1.SourceTe
 			a.readCancel()
 		}
 
+		var waitErr error
 		if a.t != nil {
 			waitErr = a.waitForRun(ctx, teardownTimeout) // wait for Run to stop running
 			if waitErr != nil {
@@ -354,12 +354,13 @@ func (a *sourcePluginAdapter) Teardown(ctx context.Context, _ cpluginv1.SourceTe
 			}
 		}
 
-		return a.impl.Teardown(ctx)
+		err := a.impl.Teardown(ctx)
+		if err == nil {
+			err = waitErr
+		}
+		return err
 	})
 
-	if err == nil {
-		err = waitErr
-	}
 	return cpluginv1.SourceTeardownResponse{}, err
 }
 
diff --git a/source_test.go b/source_test.go
index 9d1ea2e4..eacaebd0 100644
--- a/source_test.go
+++ b/source_test.go
@@ -22,12 +22,10 @@ import (
 	"testing"
 	"time"
 
-	"github.com/conduitio/conduit-connector-sdk/internal/cchan"
-
-	"github.com/conduitio/conduit-connector-sdk/internal"
-
 	"github.com/conduitio/conduit-connector-protocol/cpluginv1"
 	cpluginv1mock "github.com/conduitio/conduit-connector-protocol/cpluginv1/mock"
+	"github.com/conduitio/conduit-connector-sdk/internal"
+	"github.com/conduitio/conduit-connector-sdk/internal/cchan"
 	"github.com/golang/mock/gomock"
 	"github.com/matryer/is"
 	"github.com/rs/zerolog"