diff --git a/CHANGELOG.md b/CHANGELOG.md
index de382edb2b3ced4098348451109f171596fea7cb..e3e1c0d0b83bac4310c9bd6bd02f474f62099b2d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,7 @@ The format is based on [keep a changelog](http://keepachangelog.com) and this pr
 - Allow tournament joins before the start of the tournament active period.
 - Authoritative matches now complete their stop phase faster to avoid unnecessary processing.
 - Authoritative match join attempts now have their own bounded queue and no longer count towards the match call queue limit.
+- Lua runtime group create function now sets the correct default max size if one is not specified.
 
 ### Fixed
 - Correctly report execution mode in Lua runtime after hooks.
@@ -33,6 +34,7 @@ The format is based on [keep a changelog](http://keepachangelog.com) and this pr
 - Leaderboard and tournament runtime callbacks no longer trigger twice when time delays are observed.
 - Check group max allowed user when promoting a user.
 - Correct Lua runtime decoding of stream identifying parameters.
+- Correctly use optional parameters when they are passed to group creation operations.
 
 ## [2.2.1] - 2018-11-20
 ### Added
diff --git a/server/core_group.go b/server/core_group.go
index fd4e533680d1fa56addb6f727fbe408e0553bc10..a031a303b04c9bea648da037ebe766a270329675 100644
--- a/server/core_group.go
+++ b/server/core_group.go
@@ -63,42 +63,32 @@ func CreateGroup(ctx context.Context, logger *zap.Logger, db *sql.DB, userID uui
 	}
 
 	params := []interface{}{uuid.Must(uuid.NewV4()), creatorID, name, desc, avatarURL, state}
-	query := `
-INSERT INTO groups
-	(id, creator_id, name, description, avatar_url, state, edge_count)
-VALUES
-	($1, $2, $3, $4, $5, $6, 1)
-RETURNING id, creator_id, name, description, avatar_url, state, edge_count, lang_tag, max_count, metadata, create_time, update_time`
+	statements := []string{"$1", "$2", "$3", "$4", "$5", "$6"}
+
+	query := "INSERT INTO groups(id, creator_id, name, description, avatar_url, state"
+
+	// Add lang tag if any.
 	if lang != "" {
+		query += ", lang_tag"
 		params = append(params, lang)
-		query = `
-INSERT INTO groups
-	(id, creator_id, name, description, avatar_url, state, edge_count, lang_tag)
-VALUES
-	($1, $2, $3, $4, $5, $6, 1, $7)
-RETURNING id, creator_id, name, description, avatar_url, state, edge_count, lang_tag, max_count, metadata, create_time, update_time`
+		statements = append(statements, "$"+strconv.Itoa(len(params)))
+	}
+	// Add max count if any.
+	if maxCount > 0 {
+		query += ", max_count"
+		params = append(params, maxCount)
+		statements = append(statements, "$"+strconv.Itoa(len(params)))
+	}
+	// Add metadata if any.
+	if metadata != "" {
+		query += ", metadata"
+		params = append(params, metadata)
+		statements = append(statements, "$"+strconv.Itoa(len(params)))
 	}
 
-	// called from the client
-	if maxCount > 0 && metadata != "" {
-		params = append(params, maxCount, metadata) //no need to add 'lang' again
-
-		if lang != "" {
-			query = `
-INSERT INTO groups
-	(id, creator_id, name, description, avatar_url, state, edge_count, lang_tag, max_count, metadata)
-VALUES
-	($1, $2, $3, $4, $5, $6, 1, $7, $8, $9)
-RETURNING id, creator_id, name, description, avatar_url, state, edge_count, lang_tag, max_count, metadata, create_time, update_time`
-		} else {
-			query = `
-INSERT INTO groups
-	(id, creator_id, name, description, avatar_url, state, edge_count, max_count, metadata)
-VALUES
-	($1, $2, $3, $4, $5, $6, 1, $7, $8)
+	// Add the trailing edge count value.
+	query += `, edge_count) VALUES (` + strings.Join(statements, ",") + `,1)
 RETURNING id, creator_id, name, description, avatar_url, state, edge_count, lang_tag, max_count, metadata, create_time, update_time`
-		}
-	}
 
 	tx, err := db.BeginTx(ctx, nil)
 	if err != nil {
diff --git a/server/runtime_lua_nakama.go b/server/runtime_lua_nakama.go
index 5a28aa2d6d77fc4622b1128f39ae66faa08174c3..2f3eaa6d4de968148ef2e98398577b565dd8b62a 100644
--- a/server/runtime_lua_nakama.go
+++ b/server/runtime_lua_nakama.go
@@ -4401,7 +4401,7 @@ func (n *RuntimeLuaNakamaModule) groupCreate(l *lua.LState) int {
 		}
 		metadataStr = string(metadataBytes)
 	}
-	maxCount := l.OptInt(9, 0)
+	maxCount := l.OptInt(9, 100)
 	if maxCount < 1 || maxCount > 100 {
 		l.ArgError(9, "expects max_count to be >= 1 and <= 100")
 		return 0
diff --git a/tests/core_storage_test.go b/tests/core_storage_test.go
index c31475963d27269133eb48270b3d09a06abcc59c..3153f10844b00ac319242fd2375e3f8c7289c043 100644
--- a/tests/core_storage_test.go
+++ b/tests/core_storage_test.go
@@ -15,6 +15,7 @@
 package tests
 
 import (
+	"context"
 	"crypto/md5"
 	"fmt"
 	"testing"
@@ -40,7 +41,7 @@ func TestStorageWriteRuntimeGlobalSingle(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -55,7 +56,7 @@ func TestStorageWriteRuntimeGlobalSingle(t *testing.T) {
 			Collection: "testcollection",
 			Key:        key,
 		}}
-	readData, err := server.StorageReadObjects(logger, db, uuid.Nil, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Nil, ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, readData.Objects, "data was nil")
@@ -101,7 +102,7 @@ func TestStorageWriteRuntimeUserMultiple(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 1},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -140,7 +141,7 @@ func TestStorageWriteRuntimeGlobalSingleIfMatchNotExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, acks, "acks was not nil")
 	assert.Equal(t, codes.InvalidArgument, code, "code did not match")
@@ -160,7 +161,7 @@ func TestStorageWriteRuntimeGlobalSingleIfMatchExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -181,7 +182,7 @@ func TestStorageWriteRuntimeGlobalSingleIfMatchExists(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
 
-	acks, code, err = server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err = server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not 0")
@@ -205,7 +206,7 @@ func TestStorageWriteRuntimeGlobalSingleIfMatchExistsFail(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -226,7 +227,7 @@ func TestStorageWriteRuntimeGlobalSingleIfMatchExistsFail(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
 
-	acks, code, err = server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err = server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, acks, "acks was not nil")
 	assert.Equal(t, codes.InvalidArgument, code, "code did not match")
@@ -247,7 +248,7 @@ func TestStorageWriteRuntimeGlobalSingleIfNoneMatchNotExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, _, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, _, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -270,7 +271,7 @@ func TestStorageWriteRuntimeGlobalSingleIfNoneMatchExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, _, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, _, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -290,7 +291,7 @@ func TestStorageWriteRuntimeGlobalSingleIfNoneMatchExists(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
 
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, acks, "acks was not nil")
 	assert.Equal(t, codes.InvalidArgument, code, "code did not match")
@@ -317,7 +318,7 @@ func TestStorageWriteRuntimeGlobalMultipleIfMatchNotExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 0},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, acks, "acks was not nil")
 	assert.Equal(t, codes.InvalidArgument, code, "code did not match")
@@ -340,7 +341,7 @@ func TestStorageWritePipelineUserSingle(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, _, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, _, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -380,7 +381,7 @@ func TestStorageWritePipelineUserMultiple(t *testing.T) {
 				PermissionRead:  &wrappers.Int32Value{Value: 1},
 				PermissionWrite: &wrappers.Int32Value{Value: 1},
 			}}}
-	allAcks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	allAcks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 	if err != nil {
 		t.Fatal(err.Error())
 	}
@@ -432,7 +433,7 @@ func TestStorageWriteRuntimeGlobalMultipleSameKey(t *testing.T) {
 				PermissionRead:  &wrappers.Int32Value{Value: 1},
 				PermissionWrite: &wrappers.Int32Value{Value: 1},
 			}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -458,7 +459,7 @@ func TestStorageWriteRuntimeGlobalMultipleSameKey(t *testing.T) {
 		Key:        key,
 	}}
 
-	readData, err := server.StorageReadObjects(logger, db, uuid.Nil, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Nil, ids)
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, readData, "data was nil")
 	assert.Len(t, readData.Objects, 1, "data length was not 1")
@@ -493,7 +494,7 @@ func TestStorageWritePipelineUserMultipleSameKey(t *testing.T) {
 				PermissionRead:  &wrappers.Int32Value{Value: 1},
 				PermissionWrite: &wrappers.Int32Value{Value: 0},
 			}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not 0")
@@ -514,7 +515,7 @@ func TestStorageWritePipelineUserMultipleSameKey(t *testing.T) {
 		UserId:     uid.String(),
 	}}
 
-	readData, err := server.StorageReadObjects(logger, db, uid, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uid, ids)
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, readData, "data was nil")
 	assert.Len(t, readData.Objects, 1, "data length was not 1")
@@ -542,7 +543,7 @@ func TestStorageWritePipelineIfMatchNotExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, acks, "acks was not nil")
 	assert.Equal(t, codes.InvalidArgument, code, "code did not match")
@@ -565,7 +566,7 @@ func TestStorageWritePipelineIfMatchExistsFail(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, _, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, _, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -584,7 +585,7 @@ func TestStorageWritePipelineIfMatchExistsFail(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, acks, "acks was not nil")
 	assert.Equal(t, codes.InvalidArgument, code, "code did not match")
@@ -608,7 +609,7 @@ func TestStorageWritePipelineIfMatchExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, _, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, _, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -627,7 +628,7 @@ func TestStorageWritePipelineIfMatchExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, _, err = server.StorageWriteObjects(logger, db, false, data)
+	acks, _, err = server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -654,7 +655,7 @@ func TestStorageWritePipelineIfNoneMatchNotExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, _, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, _, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -681,7 +682,7 @@ func TestStorageWritePipelineIfNoneMatchExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, _, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, _, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -700,7 +701,7 @@ func TestStorageWritePipelineIfNoneMatchExists(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, acks, "acks was not nil")
 	assert.Equal(t, codes.InvalidArgument, code, "code did not match")
@@ -724,7 +725,7 @@ func TestStorageWritePipelinePermissionFail(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, _, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, _, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, acks, "acks was nil")
@@ -742,7 +743,7 @@ func TestStorageWritePipelinePermissionFail(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, acks, "acks was not nil")
 	assert.Equal(t, codes.InvalidArgument, code, "code did not match")
@@ -764,7 +765,7 @@ func TestStorageFetchRuntimeGlobalPrivate(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 0},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -780,7 +781,7 @@ func TestStorageFetchRuntimeGlobalPrivate(t *testing.T) {
 		Collection: "testcollection",
 		Key:        key,
 	}}
-	readData, err := server.StorageReadObjects(logger, db, uuid.Nil, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Nil, ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, data, "data was nil")
@@ -807,7 +808,7 @@ func TestStorageFetchRuntimeMixed(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 0},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -826,7 +827,7 @@ func TestStorageFetchRuntimeMixed(t *testing.T) {
 			Collection: "testcollection",
 			Key:        "notfound",
 		}}
-	readData, err := server.StorageReadObjects(logger, db, uuid.Nil, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Nil, ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -857,7 +858,7 @@ func TestStorageFetchRuntimeUserPrivate(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 0},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -874,7 +875,7 @@ func TestStorageFetchRuntimeUserPrivate(t *testing.T) {
 		Key:        key,
 		UserId:     uid.String(),
 	}}
-	readData, err := server.StorageReadObjects(logger, db, uuid.Nil, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Nil, ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -903,7 +904,7 @@ func TestStorageFetchPipelineGlobalPrivate(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 0},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -919,7 +920,7 @@ func TestStorageFetchPipelineGlobalPrivate(t *testing.T) {
 		Collection: "testcollection",
 		Key:        key,
 	}}
-	readData, err := server.StorageReadObjects(logger, db, uuid.Must(uuid.NewV4()), ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Must(uuid.NewV4()), ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -943,7 +944,7 @@ func TestStorageFetchPipelineUserPrivate(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 0},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -960,7 +961,7 @@ func TestStorageFetchPipelineUserPrivate(t *testing.T) {
 		Key:        key,
 		UserId:     uid.String(),
 	}}
-	readData, err := server.StorageReadObjects(logger, db, uid, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uid, ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -984,7 +985,7 @@ func TestStorageFetchPipelineUserRead(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 1},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1001,7 +1002,7 @@ func TestStorageFetchPipelineUserRead(t *testing.T) {
 		Key:        key,
 		UserId:     uid.String(),
 	}}
-	readData, err := server.StorageReadObjects(logger, db, uid, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uid, ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1032,7 +1033,7 @@ func TestStorageFetchPipelineUserPublic(t *testing.T) {
 			PermissionRead:  &wrappers.Int32Value{Value: 2},
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1048,7 +1049,7 @@ func TestStorageFetchPipelineUserPublic(t *testing.T) {
 		Key:        key,
 		UserId:     uid.String(),
 	}}
-	dataRead, err := server.StorageReadObjects(logger, db, uid, ids)
+	dataRead, err := server.StorageReadObjects(context.Background(), logger, db, uid, ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1080,7 +1081,7 @@ func TestStorageFetchPipelineUserOtherRead(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1097,7 +1098,7 @@ func TestStorageFetchPipelineUserOtherRead(t *testing.T) {
 		Key:        key,
 		UserId:     uid.String(),
 	}}
-	readData, err := server.StorageReadObjects(logger, db, uuid.Must(uuid.NewV4()), ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Must(uuid.NewV4()), ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, readData, "data was nil")
@@ -1121,7 +1122,7 @@ func TestStorageFetchPipelineUserOtherPublic(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1139,7 +1140,7 @@ func TestStorageFetchPipelineUserOtherPublic(t *testing.T) {
 		Key:        key,
 		UserId:     uid.String(),
 	}}
-	readData, err := server.StorageReadObjects(logger, db, uuid.Must(uuid.NewV4()), ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Must(uuid.NewV4()), ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, readData, "data was nil")
@@ -1178,7 +1179,7 @@ func TestStorageFetchPipelineUserOtherPublicMixed(t *testing.T) {
 				PermissionWrite: &wrappers.Int32Value{Value: 0},
 			},
 		}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1204,7 +1205,7 @@ func TestStorageFetchPipelineUserOtherPublicMixed(t *testing.T) {
 			UserId:     uid.String(),
 		},
 	}
-	readData, err := server.StorageReadObjects(logger, db, uuid.Must(uuid.NewV4()), ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Must(uuid.NewV4()), ids)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.NotNil(t, readData, "data was nil")
@@ -1233,7 +1234,7 @@ func TestStorageRemoveRuntimeGlobalPublic(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1252,7 +1253,7 @@ func TestStorageRemoveRuntimeGlobalPublic(t *testing.T) {
 		}},
 	}
 
-	_, err = server.StorageDeleteObjects(logger, db, true, deleteIDs)
+	_, err = server.StorageDeleteObjects(context.Background(), logger, db, true, deleteIDs)
 	assert.Nil(t, err, "err was not nil")
 }
 
@@ -1271,7 +1272,7 @@ func TestStorageRemoveRuntimeGlobalPrivate(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1289,7 +1290,7 @@ func TestStorageRemoveRuntimeGlobalPrivate(t *testing.T) {
 			Key:        key,
 		}},
 	}
-	_, err = server.StorageDeleteObjects(logger, db, true, deleteIDs)
+	_, err = server.StorageDeleteObjects(context.Background(), logger, db, true, deleteIDs)
 	assert.Nil(t, err, "err was not nil")
 }
 
@@ -1310,7 +1311,7 @@ func TestStorageRemoveRuntimeUserPublic(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1328,7 +1329,7 @@ func TestStorageRemoveRuntimeUserPublic(t *testing.T) {
 			Key:        key,
 		}},
 	}
-	_, err = server.StorageDeleteObjects(logger, db, true, deleteIDs)
+	_, err = server.StorageDeleteObjects(context.Background(), logger, db, true, deleteIDs)
 	assert.Nil(t, err, "err was not nil")
 }
 
@@ -1349,7 +1350,7 @@ func TestStorageRemoveRuntimeUserPrivate(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1367,7 +1368,7 @@ func TestStorageRemoveRuntimeUserPrivate(t *testing.T) {
 			Key:        key,
 		}},
 	}
-	_, err = server.StorageDeleteObjects(logger, db, true, deleteIDs)
+	_, err = server.StorageDeleteObjects(context.Background(), logger, db, true, deleteIDs)
 	assert.Nil(t, err, "err was not nil")
 
 	ids := []*api.ReadStorageObjectId{{
@@ -1376,7 +1377,7 @@ func TestStorageRemoveRuntimeUserPrivate(t *testing.T) {
 		UserId:     uid.String(),
 	}}
 
-	readData, err := server.StorageReadObjects(logger, db, uuid.Nil, ids)
+	readData, err := server.StorageReadObjects(context.Background(), logger, db, uuid.Nil, ids)
 	assert.Nil(t, err, "err was not nil")
 	assert.Len(t, readData.Objects, 0, "data length was not 0")
 }
@@ -1398,7 +1399,7 @@ func TestStorageRemovePipelineUserWrite(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1416,7 +1417,7 @@ func TestStorageRemovePipelineUserWrite(t *testing.T) {
 			Key:        key,
 		}},
 	}
-	_, err = server.StorageDeleteObjects(logger, db, true, deleteIDs)
+	_, err = server.StorageDeleteObjects(context.Background(), logger, db, true, deleteIDs)
 	assert.Nil(t, err, "err was not nil")
 }
 
@@ -1437,7 +1438,7 @@ func TestStorageRemovePipelineUserDenied(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 0},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1455,7 +1456,7 @@ func TestStorageRemovePipelineUserDenied(t *testing.T) {
 			Key:        key,
 		}},
 	}
-	code, err = server.StorageDeleteObjects(logger, db, false, deleteIDs)
+	code, err = server.StorageDeleteObjects(context.Background(), logger, db, false, deleteIDs)
 	assert.NotNil(t, err, "err was nil")
 	assert.Equal(t, code, codes.InvalidArgument, "code did not match InvalidArgument.")
 }
@@ -1471,7 +1472,7 @@ func TestStorageRemoveRuntimeGlobalIfMatchNotExists(t *testing.T) {
 			Version:    "fail",
 		}},
 	}
-	code, err := server.StorageDeleteObjects(logger, db, true, deleteIDs)
+	code, err := server.StorageDeleteObjects(context.Background(), logger, db, true, deleteIDs)
 	assert.NotNil(t, err, "err was nil")
 	assert.Equal(t, code, codes.InvalidArgument, "code did not match InvalidArgument.")
 }
@@ -1491,7 +1492,7 @@ func TestStorageRemoveRuntimeGlobalIfMatchRejected(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1510,7 +1511,7 @@ func TestStorageRemoveRuntimeGlobalIfMatchRejected(t *testing.T) {
 			Version:    "fail",
 		}},
 	}
-	code, err = server.StorageDeleteObjects(logger, db, true, deleteIDs)
+	code, err = server.StorageDeleteObjects(context.Background(), logger, db, true, deleteIDs)
 	assert.NotNil(t, err, "err was not nil")
 	assert.Equal(t, code, codes.InvalidArgument, "code did not match InvalidArgument.")
 }
@@ -1530,7 +1531,7 @@ func TestStorageRemoveRuntimeGlobalIfMatch(t *testing.T) {
 			PermissionWrite: &wrappers.Int32Value{Value: 1},
 		}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, true, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, true, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1549,7 +1550,7 @@ func TestStorageRemoveRuntimeGlobalIfMatch(t *testing.T) {
 			Version:    acks.Acks[0].Version,
 		}},
 	}
-	code, err = server.StorageDeleteObjects(logger, db, true, deleteIDs)
+	code, err = server.StorageDeleteObjects(context.Background(), logger, db, true, deleteIDs)
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, code, codes.OK, "code did not match OK.")
 }
@@ -1584,14 +1585,14 @@ func TestStorageListRuntimeUser(t *testing.T) {
 				PermissionWrite: &wrappers.Int32Value{Value: 0},
 			},
 		}}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
 	assert.NotNil(t, acks, "acks was nil")
 	assert.Len(t, acks.Acks, 3, "acks length was not 3")
 
-	list, code, err := server.StorageListObjects(logger, db, uuid.Nil, uid, "testcollection", 10, "")
+	list, code, err := server.StorageListObjects(context.Background(), logger, db, uuid.Nil, uid, "testcollection", 10, "")
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1631,14 +1632,14 @@ func TestStorageListPipelineUserSelf(t *testing.T) {
 				PermissionWrite: &wrappers.Int32Value{Value: 0},
 			}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
 	assert.NotNil(t, acks, "acks was nil")
 	assert.Len(t, acks.Acks, 3, "acks length was not 3")
 
-	list, code, err := server.StorageListObjects(logger, db, uid, uid, collection, 10, "")
+	list, code, err := server.StorageListObjects(context.Background(), logger, db, uid, uid, collection, 10, "")
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
@@ -1680,14 +1681,14 @@ func TestStorageListPipelineUserOther(t *testing.T) {
 				PermissionWrite: &wrappers.Int32Value{Value: 0},
 			}},
 	}
-	acks, code, err := server.StorageWriteObjects(logger, db, false, data)
+	acks, code, err := server.StorageWriteObjects(context.Background(), logger, db, false, data)
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
 	assert.NotNil(t, acks, "acks was nil")
 	assert.Len(t, acks.Acks, 3, "acks length was not 3")
 
-	values, code, err := server.StorageListObjects(logger, db, uuid.Must(uuid.NewV4()), uid, collection, 10, "")
+	values, code, err := server.StorageListObjects(context.Background(), logger, db, uuid.Must(uuid.NewV4()), uid, collection, 10, "")
 
 	assert.Nil(t, err, "err was not nil")
 	assert.Equal(t, codes.OK, code, "code was not OK")
diff --git a/tests/runtime_test.go b/tests/runtime_test.go
index 116b6636b5134892a907696310758ce44eedfa98..03b6f62f7a51bef128717107d049c7e5cdc21a41 100644
--- a/tests/runtime_test.go
+++ b/tests/runtime_test.go
@@ -15,6 +15,7 @@
 package tests
 
 import (
+	"context"
 	"errors"
 	"io/ioutil"
 	"net/http"
@@ -320,7 +321,7 @@ nakama.register_rpc(test.printWorld, "helloworld")`,
 	}
 
 	payload := "Hello World"
-	result, err, _ := fn(nil, "", "", 0, "", "", "", payload)
+	result, err, _ := fn(context.Background(), nil, "", "", 0, "", "", "", payload)
 	if err != nil {
 		t.Fatal(err.Error())
 	}
@@ -398,7 +399,7 @@ nakama.register_rpc(test, "test")`,
 		t.Fatal("Expected RPC function to be registered")
 	}
 
-	result, err, _ := fn(nil, "", "", 0, "", "", "", "")
+	result, err, _ := fn(context.Background(), nil, "", "", 0, "", "", "", "")
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -429,7 +430,7 @@ nakama.register_rpc(test, "test")`,
 	}
 
 	payload := "{\"key\":\"value\"}"
-	result, err, _ := fn(nil, "", "", 0, "", "", "", payload)
+	result, err, _ := fn(context.Background(), nil, "", "", 0, "", "", "", payload)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -460,7 +461,7 @@ nakama.register_rpc(test, "test")`,
 	}
 
 	payload := "{\"key\":\"value\"}"
-	result, err, _ := fn(nil, "", "", 0, "", "", "", payload)
+	result, err, _ := fn(context.Background(), nil, "", "", 0, "", "", "", payload)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -491,7 +492,7 @@ nakama.register_rpc(test, "test")`,
 	}
 
 	payload := "{\"key\":\"value\"}"
-	result, err, _ := fn(nil, "", "", 0, "", "", "", payload)
+	result, err, _ := fn(context.Background(), nil, "", "", 0, "", "", "", payload)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -522,7 +523,7 @@ nakama.register_rpc(test, "test")`,
 	}
 
 	payload := "{\"key\":\"value\"}"
-	result, err, _ := fn(nil, "", "", 0, "", "", "", payload)
+	result, err, _ := fn(context.Background(), nil, "", "", 0, "", "", "", payload)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -579,7 +580,7 @@ nakama.register_rpc(test, "test")`,
 	}
 
 	payload := "{\"key\":\"value\"}"
-	result, err, _ := fn(nil, "", "", 0, "", "", "", payload)
+	result, err, _ := fn(context.Background(), nil, "", "", 0, "", "", "", payload)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -612,7 +613,7 @@ nakama.register_rpc(test, "test")`,
 
 	payload := "something_to_encrypt"
 	hash, _ := bcrypt.GenerateFromPassword([]byte(payload), bcrypt.DefaultCost)
-	result, err, _ := fn(nil, "", "", 0, "", "", "", string(hash))
+	result, err, _ := fn(context.Background(), nil, "", "", 0, "", "", "", string(hash))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -980,7 +981,7 @@ local users = nk.group_users_list(group.id)
 for i, u in ipairs(users)
 do
   assert(u.user.id == user_id, "'u.id' must be equal to user_id")
-	assert(u.state == 1, "'u.state' must be equal to 1 / superadmin")
+	assert(u.state == 0, "'u.state' must be equal to 0 / superadmin")
 end
 
 local groups = nk.user_groups_list(user_id)
@@ -988,7 +989,7 @@ for i, g in ipairs(groups)
 do
 	print(nk.json_encode(g))
   assert(g.group.name == group_update_name, "'g.name' must be equal to group_update_name")
-	assert(g.state == 1, "'g.state' must be equal to 1 / superadmin")
+	assert(g.state == 0, "'g.state' must be equal to 0 / superadmin")
 end
 
 nk.group_delete(group.id)`,
diff --git a/tests/util.go b/tests/util.go
index d289c076477847ee2fe879e24812b951e06425eb..112d9947f1dc35814f6c99d73228d8994a6923b0 100644
--- a/tests/util.go
+++ b/tests/util.go
@@ -23,6 +23,7 @@ import (
 	"github.com/gofrs/uuid"
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/heroiclabs/nakama/api"
+	"github.com/heroiclabs/nakama/apigrpc"
 	"github.com/heroiclabs/nakama/rtapi"
 	"github.com/heroiclabs/nakama/server"
 	"go.uber.org/zap"
@@ -86,6 +87,9 @@ func (d *DummySession) ClientIP() string {
 func (d *DummySession) ClientPort() string {
 	return ""
 }
+func (d *DummySession) Context() context.Context {
+	return context.Background()
+}
 func (d *DummySession) Send(isStream bool, mode uint8, envelope *rtapi.Envelope) error {
 	d.messages = append(d.messages, envelope)
 	return nil
@@ -155,11 +159,11 @@ func NewAPIServer(t *testing.T, runtime *server.Runtime) (*server.ApiServer, *se
 	router := &DummyMessageRouter{}
 	tracker := &server.LocalTracker{}
 	pipeline := server.NewPipeline(logger, config, db, jsonpbMarshaler, jsonpbUnmarshaler, nil, nil, nil, tracker, router, runtime)
-	apiServer := server.StartApiServer(logger, logger, db, jsonpbMarshaler, jsonpbUnmarshaler, config, nil, nil, nil, nil, nil, tracker, router, pipeline, runtime)
+	apiServer := server.StartApiServer(logger, logger, db, jsonpbMarshaler, jsonpbUnmarshaler, config, nil, nil, nil, nil, nil, nil, tracker, router, pipeline, runtime)
 	return apiServer, pipeline
 }
 
-func NewSession(t *testing.T, customID string) (*grpc.ClientConn, api.NakamaClient, *api.Session, context.Context) {
+func NewSession(t *testing.T, customID string) (*grpc.ClientConn, apigrpc.NakamaClient, *api.Session, context.Context) {
 	ctx := context.Background()
 	outgoingCtx := metadata.NewOutgoingContext(ctx, metadata.New(map[string]string{
 		"authorization": "Basic " + base64.StdEncoding.EncodeToString([]byte("defaultkey:")),
@@ -169,7 +173,7 @@ func NewSession(t *testing.T, customID string) (*grpc.ClientConn, api.NakamaClie
 		t.Fatal(err)
 	}
 
-	client := api.NewNakamaClient(conn)
+	client := apigrpc.NewNakamaClient(conn)
 	session, err := client.AuthenticateCustom(outgoingCtx, &api.AuthenticateCustomRequest{
 		Account: &api.AccountCustom{
 			Id: customID,
@@ -183,7 +187,7 @@ func NewSession(t *testing.T, customID string) (*grpc.ClientConn, api.NakamaClie
 	return conn, client, session, outgoingCtx
 }
 
-func NewAuthenticatedAPIClient(t *testing.T, customID string) (*grpc.ClientConn, api.NakamaClient, *api.Session, context.Context) {
+func NewAuthenticatedAPIClient(t *testing.T, customID string) (*grpc.ClientConn, apigrpc.NakamaClient, *api.Session, context.Context) {
 	conn, _, session, _ := NewSession(t, customID)
 	conn.Close()
 
@@ -196,7 +200,7 @@ func NewAuthenticatedAPIClient(t *testing.T, customID string) (*grpc.ClientConn,
 		t.Fatal(err)
 	}
 
-	client := api.NewNakamaClient(conn)
+	client := apigrpc.NewNakamaClient(conn)
 	return conn, client, session, outgoingCtx
 }