[tests] Simplify Cassandra e2e test cleanup (#4794)

Follow-up on #4776

Signed-off-by: Yuri Shkuro <[email protected]>

plugin/storage/integration/cassandra_test.go

@@ -37,28 +37,25 @@ var errInitializeCassandraDependencyWriter = errors.New("failed to initialize ca
 type CassandraStorageIntegration struct {
 	StorageIntegration
 
-	logger *zap.Logger
+	session dbsession.Session
+	logger  *zap.Logger
 }
 
 func newCassandraStorageIntegration() *CassandraStorageIntegration {
-	return &CassandraStorageIntegration{
+	s := &CassandraStorageIntegration{
 		StorageIntegration: StorageIntegration{
 			GetDependenciesReturnsSource: true,
 
 			Refresh:  func() error { return nil },
-			CleanUp:  func() error { return nil },
 			SkipList: []string{"Tags_+_Operation_name_+_Duration_range", "Tags_+_Duration_range", "Tags_+_Operation_name_+_max_Duration", "Tags_+_max_Duration", "Operation_name_+_Duration_range", "Duration_range", "max_Duration"},
 		},
 	}
+	s.CleanUp = s.cleanUp
+	return s
 }
 
-func (s *CassandraStorageIntegration) cleanUp(q dbsession.Session) (func() error, error) {
-	return func() error {
-		if err := q.Query("TRUNCATE traces").Exec(); err != nil {
-			return err
-		}
-		return nil
-	}, nil
+func (s *CassandraStorageIntegration) cleanUp() error {
+	return s.session.Query("TRUNCATE traces").Exec()
 }
 
 func (s *CassandraStorageIntegration) initializeCassandraFactory(flags []string) (*cassandra.Factory, error) {
@@ -82,15 +79,13 @@ func (s *CassandraStorageIntegration) initializeCassandra() error {
 	if err != nil {
 		return err
 	}
+	s.session = f.PrimarySession()
 	if s.SpanWriter, err = f.CreateSpanWriter(); err != nil {
 		return err
 	}
 	if s.SpanReader, err = f.CreateSpanReader(); err != nil {
 		return err
 	}
-	if s.CleanUp, err = s.cleanUp(f.PrimarySession()); err != nil {
-		return err
-	}
 	if err = s.initializeDependencyReaderAndWriter(f); err != nil {
 		return err
 	}

plugin/storage/integration/integration.go

@@ -60,6 +60,7 @@ type StorageIntegration struct {
 
 	// List of tests which has to be skipped, it can be regex too.
 	SkipList []string
+
 	// CleanUp() should ensure that the storage backend is clean before another test.
 	// called either before or after each test, and should be idempotent
 	CleanUp func() error