diff --git a/scylla_shard_aware_extension_integration_test.go b/scylla_shard_aware_extension_integration_test.go
index ea8b27f1f..d9f10e03c 100644
--- a/scylla_shard_aware_extension_integration_test.go
+++ b/scylla_shard_aware_extension_integration_test.go
@@ -25,11 +25,20 @@ func TestSession_GetShardAwareRoutingInfo_Integration(t *testing.T) {
 		if err != nil {
 			t.Fatalf("failed to create session '%v'", err)
 		}
-		defer session.Close()
+
+		// best practice: add clean up
+		t.Cleanup(func() {
+			defer session.Close() // close session after tests
+
+			// clear DB
+			if err := createTable(session, `DROP KEYSPACE IF EXISTS `+keyspace); err != nil {
+				t.Logf(fmt.Sprintf("unable to drop keyspace: %v", err))
+			}
+		})
 
 		err = createTable(session, `DROP KEYSPACE IF EXISTS `+keyspace)
 		if err != nil {
-			panic(fmt.Sprintf("unable to drop keyspace: %v", err))
+			t.Fatalf(fmt.Sprintf("unable to drop keyspace: %v", err))
 		}
 
 		err = createTable(session, fmt.Sprintf(`CREATE KEYSPACE %s
@@ -37,8 +46,12 @@ func TestSession_GetShardAwareRoutingInfo_Integration(t *testing.T) {
 		'class': 'NetworkTopologyStrategy',
 		'replication_factor' : %d
 	}`, keyspace, *flagRF))
+		if err != nil {
+			t.Fatalf(fmt.Sprintf("unable to create keyspace: %v", err))
+		}
 
-		if err := createTable(session, fmt.Sprintf("CREATE TABLE %s.%s (first_id int, second_id int, third_id int, PRIMARY KEY ((first_id, second_id)))", keyspace, table)); err != nil {
+		err = createTable(session, fmt.Sprintf("CREATE TABLE %s.%s (first_id int, second_id int, third_id int, PRIMARY KEY ((first_id, second_id)))", keyspace, table))
+		if err != nil {
 			t.Fatalf("failed to create table with error '%v'", err)
 		}
 	}
@@ -52,6 +65,7 @@ func TestSession_GetShardAwareRoutingInfo_Integration(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to create session '%v'", err)
 	}
+	defer session.Close()
 
 	info, err := session.GetShardAwareRoutingInfo(table, []string{"first_id", "second_id", "third_id"}, 1, 2, 3)
 	if err != nil {