diff --git a/tests/php/Store/SessionStoreTest.php b/tests/php/Store/SessionStoreTest.php
index af54558..443df17 100644
--- a/tests/php/Store/SessionStoreTest.php
+++ b/tests/php/Store/SessionStoreTest.php
@@ -91,7 +91,7 @@ public function preparedQuery($sql, $parameters, $errorLevel = E_USER_ERROR)
                 $this->preparedQueryNum++;
             }
         };
-        DB::set_conn($mock);
+        DB::set_conn($mock, DB::CONN_PRIMARY);
 
         // Replicate the deserialisation that happens on session start
         $store->__unserialize($serialised);
@@ -100,6 +100,6 @@ public function preparedQuery($sql, $parameters, $errorLevel = E_USER_ERROR)
         $this->assertSame(0, $mock->preparedQueryNum);
 
         // Finish the test and allow mock assertions to fail the test
-        DB::set_conn($connection);
+        DB::set_conn($connection, DB::CONN_PRIMARY);
     }
 }