hachikuji commented on a change in pull request #9512:
URL: https://github.com/apache/kafka/pull/9512#discussion_r537757487



##########
File path: raft/src/test/java/org/apache/kafka/raft/MockLog.java
##########
@@ -472,4 +490,106 @@ private EpochStartOffset(int epoch, long startOffset) {
         }
     }
 
+    final class MockRawSnapshotWriter implements RawSnapshotWriter {
+        private final OffsetAndEpoch snapshotId;
+        private ByteBufferOutputStream data;
+        private boolean frozen;
+
+        public MockRawSnapshotWriter(OffsetAndEpoch snapshotId) {
+            this.snapshotId = snapshotId;
+            this.data = new ByteBufferOutputStream(0);
+            this.frozen = false;
+        }
+
+        @Override
+        public OffsetAndEpoch snapshotId() {
+            return snapshotId;
+        }
+
+        @Override
+        public long sizeInBytes() {
+            return data.position();
+        }
+
+        @Override
+        public void append(ByteBuffer buffer) {
+            if (frozen) {
+                throw new RuntimeException("Snapshot is already frozen " + 
snapshotId);
+            }
+
+            data.write(buffer);
+        }
+
+        @Override
+        public boolean isFrozen() {
+            return frozen;
+        }
+
+        @Override
+        public void freeze() {
+            if (frozen) {
+                throw new RuntimeException("Snapshot is already frozen " + 
snapshotId);
+            }
+
+            frozen = true;
+            ByteBuffer buffer = data.buffer();
+            buffer.flip();
+
+            snapshots.putIfAbsent(snapshotId, new 
MockRawSnapshotReader(snapshotId, buffer));
+        }
+
+        @Override
+        public void close() {}
+    }
+
+    final static class MockRawSnapshotReader implements RawSnapshotReader {
+        private final OffsetAndEpoch snapshotId;
+        private final MemoryRecords data;
+
+        MockRawSnapshotReader(OffsetAndEpoch snapshotId, ByteBuffer data) {
+            this.snapshotId = snapshotId;
+            this.data = MemoryRecords.readableRecords(data);
+        }
+
+        @Override
+        public OffsetAndEpoch snapshotId() {
+            return snapshotId;
+        }
+
+        @Override
+        public long sizeInBytes() {
+            return data.sizeInBytes();
+        }
+
+        @Override
+        public Iterator<RecordBatch> iterator() {
+            return new Iterator<RecordBatch>() {

Review comment:
       nit: we cab use `covariantCast`?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to