diff --git a/keyExchange/utils_test.go b/keyExchange/utils_test.go
index c11728b1afd66ace0983fce8dcf2134978f35150..bda05f0b822888261f8f61f460840536dd9ddbc5 100644
--- a/keyExchange/utils_test.go
+++ b/keyExchange/utils_test.go
@@ -101,11 +101,7 @@ func (t *testNetworkManagerGeneric) GetStoppable() stoppable.Stoppable {
 
 func InitTestingContextGeneric(i interface{}) (*storage.Session, interfaces.NetworkManager) {
 	switch i.(type) {
-	case *testing.T:
-		break
-	case *testing.M:
-		break
-	case *testing.B:
+	case *testing.T, *testing.M, *testing.B, *testing.PB:
 		break
 	default:
 		globals.Log.FATAL.Panicf("InitTestingSession is restricted to testing only. Got %T", i)
@@ -208,11 +204,7 @@ func (t *testNetworkManagerFullExchange) GetStoppable() stoppable.Stoppable {
 
 func InitTestingContextFullExchange(i interface{}) (*storage.Session, *switchboard.Switchboard, interfaces.NetworkManager) {
 	switch i.(type) {
-	case *testing.T:
-		break
-	case *testing.M:
-		break
-	case *testing.B:
+	case *testing.T, *testing.M, *testing.B, *testing.PB:
 		break
 	default:
 		globals.Log.FATAL.Panicf("InitTestingSession is restricted to testing only. Got %T", i)
diff --git a/storage/e2e/session.go b/storage/e2e/session.go
index 43eacbf80f676c1eeefbab5adcb395b31cbcd08a..78997dc00d7e78ca1b5badb5d82ef6e676837f7a 100644
--- a/storage/e2e/session.go
+++ b/storage/e2e/session.go
@@ -243,11 +243,7 @@ func getSessionIDFromBaseKey(baseKey *cyclic.Int) SessionID {
 // FOR TESTING PURPOSES ONLY
 func GetSessionIDFromBaseKeyForTesting(baseKey *cyclic.Int, i interface{}) SessionID {
 	switch i.(type) {
-	case *testing.T:
-		break
-	case *testing.M:
-		break
-	case *testing.B:
+	case *testing.T, *testing.M, *testing.B, *testing.PB:
 		break
 	default:
 		globals.Log.FATAL.Panicf("GetSessionIDFromBaseKeyForTesting is restricted to testing only. Got %T", i)
diff --git a/storage/session.go b/storage/session.go
index 8f276039df5eae534421f57008921d2e90852749..ab8d23d99ef05f114fa3f17b5b0333bf2906d682 100644
--- a/storage/session.go
+++ b/storage/session.go
@@ -290,11 +290,7 @@ func (s *Session) Delete(key string) error {
 // FOR TESTING ONLY
 func InitTestingSession(i interface{}) *Session {
 	switch i.(type) {
-	case *testing.T:
-		break
-	case *testing.M:
-		break
-	case *testing.B:
+	case *testing.T, *testing.M, *testing.B, *testing.PB:
 		break
 	default:
 		globals.Log.FATAL.Panicf("InitTestingSession is restricted to testing only. Got %T", i)