@@ -52,8 +52,8 @@ func TestFirewallDBMigration(t *testing.T) {
5252 t .Skipf ("Skipping Firewall DB migration test for kvdb build" )
5353 }
5454
55- makeSQLDB := func (t * testing.T , sessionsStore session. Store ) ( * SQLDB ,
56- * SQLQueriesExecutor [SQLQueries ]) {
55+ makeSQLDB := func (t * testing.T ,
56+ sessionsStore session. Store ) * SQLQueriesExecutor [SQLQueries ] {
5757
5858 testDBStore := NewTestDBWithSessions (t , sessionsStore , clock )
5959
@@ -64,13 +64,13 @@ func TestFirewallDBMigration(t *testing.T) {
6464
6565 queries := sqlc .NewForType (baseDB , baseDB .BackendType )
6666
67- return store , NewSQLQueriesExecutor (baseDB , queries )
67+ return NewSQLQueriesExecutor (baseDB , queries )
6868 }
6969
7070 // The assertMigrationResults function will currently assert that
7171 // the migrated kv stores entries in the SQLDB match the original kv
7272 // stores entries in the BoltDB.
73- assertMigrationResults := func (t * testing.T , store * SQLDB ,
73+ assertMigrationResults := func (t * testing.T , store SQLQueries ,
7474 kvEntries []* kvEntry ) {
7575
7676 var (
@@ -83,9 +83,7 @@ func TestFirewallDBMigration(t *testing.T) {
8383 getRuleID := func (ruleName string ) int64 {
8484 ruleID , ok := ruleIDs [ruleName ]
8585 if ! ok {
86- ruleID , err = store .db .GetRuleID (
87- ctx , ruleName ,
88- )
86+ ruleID , err = store .GetRuleID (ctx , ruleName )
8987 require .NoError (t , err )
9088
9189 ruleIDs [ruleName ] = ruleID
@@ -97,7 +95,7 @@ func TestFirewallDBMigration(t *testing.T) {
9795 getGroupID := func (groupAlias []byte ) int64 {
9896 groupID , ok := groupIDs [string (groupAlias )]
9997 if ! ok {
100- groupID , err = store .db . GetSessionIDByAlias (
98+ groupID , err = store .GetSessionIDByAlias (
10199 ctx , groupAlias ,
102100 )
103101 require .NoError (t , err )
@@ -111,7 +109,7 @@ func TestFirewallDBMigration(t *testing.T) {
111109 getFeatureID := func (featureName string ) int64 {
112110 featureID , ok := featureIDs [featureName ]
113111 if ! ok {
114- featureID , err = store .db . GetFeatureID (
112+ featureID , err = store .GetFeatureID (
115113 ctx , featureName ,
116114 )
117115 require .NoError (t , err )
@@ -125,7 +123,7 @@ func TestFirewallDBMigration(t *testing.T) {
125123 // First we extract all migrated kv entries from the SQLDB,
126124 // in order to be able to compare them to the original kv
127125 // entries, to ensure that the migration was successful.
128- sqlKvEntries , err := store .db . ListAllKVStoresRecords (ctx )
126+ sqlKvEntries , err := store .ListAllKVStoresRecords (ctx )
129127 require .NoError (t , err )
130128 require .Equal (t , len (kvEntries ), len (sqlKvEntries ))
131129
@@ -141,7 +139,7 @@ func TestFirewallDBMigration(t *testing.T) {
141139 ruleID := getRuleID (entry .ruleName )
142140
143141 if entry .groupAlias .IsNone () {
144- sqlVal , err := store .db . GetGlobalKVStoreRecord (
142+ sqlVal , err := store .GetGlobalKVStoreRecord (
145143 ctx ,
146144 sqlc.GetGlobalKVStoreRecordParams {
147145 Key : entry .key ,
@@ -159,7 +157,7 @@ func TestFirewallDBMigration(t *testing.T) {
159157 groupAlias := entry .groupAlias .UnwrapOrFail (t )
160158 groupID := getGroupID (groupAlias [:])
161159
162- v , err := store .db . GetGroupKVStoreRecord (
160+ v , err := store .GetGroupKVStoreRecord (
163161 ctx ,
164162 sqlc.GetGroupKVStoreRecordParams {
165163 Key : entry .key ,
@@ -184,7 +182,7 @@ func TestFirewallDBMigration(t *testing.T) {
184182 entry .featureName .UnwrapOrFail (t ),
185183 )
186184
187- sqlVal , err := store .db . GetFeatureKVStoreRecord (
185+ sqlVal , err := store .GetFeatureKVStoreRecord (
188186 ctx ,
189187 sqlc.GetFeatureKVStoreRecordParams {
190188 Key : entry .key ,
@@ -290,7 +288,7 @@ func TestFirewallDBMigration(t *testing.T) {
290288
291289 // Create the SQL store that we will migrate the data
292290 // to.
293- sqlStore , txEx := makeSQLDB (t , sessionsStore )
291+ txEx := makeSQLDB (t , sessionsStore )
294292
295293 // Perform the migration.
296294 //
@@ -299,15 +297,20 @@ func TestFirewallDBMigration(t *testing.T) {
299297 var opts sqldb.MigrationTxOptions
300298 err = txEx .ExecTx (ctx , & opts ,
301299 func (tx SQLQueries ) error {
302- return MigrateFirewallDBToSQL (
300+ err = MigrateFirewallDBToSQL (
303301 ctx , firewallStore .DB , tx ,
304302 )
303+ if err != nil {
304+ return err
305+ }
306+
307+ // Assert migration results.
308+ assertMigrationResults (t , tx , entries )
309+
310+ return nil
305311 }, sqldb .NoOpReset ,
306312 )
307313 require .NoError (t , err )
308-
309- // Assert migration results.
310- assertMigrationResults (t , sqlStore , entries )
311314 })
312315 }
313316}
0 commit comments