@@ -156,7 +156,7 @@ protected Message newState(ProjectId projectId) {
156156 private EntityRecordWithColumns newRecordWithColumns (RecordStorage <ProjectId > storage ) {
157157 EntityRecord record = newStorageRecord ();
158158 Entity <ProjectId , Project > entity = new TestCounterEntity (newId ());
159- EntityRecordWithColumns recordWithColumns = create (record , entity , storage . entityClass () );
159+ EntityRecordWithColumns recordWithColumns = create (record , entity , storage );
160160 return recordWithColumns ;
161161 }
162162
@@ -220,7 +220,7 @@ void testPersistColumns() {
220220 .setVersion (versionValue )
221221 .vBuild ();
222222 DsRecordStorage <ProjectId > storage = newStorage (TestCounterEntity .class );
223- EntityRecordWithColumns recordWithColumns = create (record , entity , storage . entityClass () );
223+ EntityRecordWithColumns recordWithColumns = create (record , entity , storage );
224224 ImmutableSet <String > columns = recordWithColumns .columnNames ()
225225 .stream ()
226226 .map (ColumnName ::value )
@@ -318,7 +318,7 @@ void testLifecycleFlags() {
318318 .vBuild ();
319319 TestCounterEntity entity = new TestCounterEntity (id );
320320 RecordStorage <ProjectId > storage = newStorage (TestCounterEntity .class );
321- EntityRecordWithColumns recordWithColumns = create (record , entity , storage . entityClass () );
321+ EntityRecordWithColumns recordWithColumns = create (record , entity , storage );
322322 storage .write (id , recordWithColumns );
323323
324324 RecordReadRequest <ProjectId > request = new RecordReadRequest <>(id );
@@ -597,7 +597,7 @@ void testQueryEntityWithoutLifecycleById() {
597597 ProjectId id = newId ();
598598 EntityRecord record = newEntityRecord (id , newState (id ));
599599 EntityWithoutLifecycle entity = new EntityWithoutLifecycle (id );
600- storage .writeRecord (entity .id (), create (record , entity , storage . entityClass () ));
600+ storage .writeRecord (entity .id (), create (record , entity , storage ));
601601
602602 // Create ID filter.
603603 List <Any > targetIds = singletonList (pack (entity .id ()));
0 commit comments