-
Notifications
You must be signed in to change notification settings - Fork 272
/
RecordAccumulatorTest.java
140 lines (119 loc) · 4.83 KB
/
RecordAccumulatorTest.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
// Copyright 2017 The Nomulus Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package google.registry.tools;
import static com.google.common.truth.Truth.assertThat;
import static google.registry.tools.LevelDbLogReader.BLOCK_SIZE;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.EntityTranslator;
import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableSet;
import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
import google.registry.testing.AppEngineRule;
import google.registry.tools.LevelDbLogReader.ChunkType;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
@RunWith(JUnit4.class)
public class RecordAccumulatorTest {
private static final int BASE_ID = 1001;
private static final String TEST_ENTITY_KIND = "TestEntity";
@Rule public final TemporaryFolder tempFs = new TemporaryFolder();
@Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build();
@Test
public void testReadDirectory() throws IOException {
File subdir = tempFs.newFolder("folder");
LevelDbFileBuilder builder = new LevelDbFileBuilder(new File(subdir, "data1"));
// Note that we need to specify property values as "Long" for property comparisons to work
// correctly because that's how they are deserialized from protos.
ComparableEntity e1 =
builder.addEntityProto(
BASE_ID,
Property.create("eeny", 100L),
Property.create("meeny", 200L),
Property.create("miney", 300L));
ComparableEntity e2 =
builder.addEntityProto(
BASE_ID + 1,
Property.create("eeny", 100L),
Property.create("meeny", 200L),
Property.create("miney", 300L));
builder.build();
builder = new LevelDbFileBuilder(new File(subdir, "data2"));
// Duplicate of the record in the other file.
builder.addEntityProto(
BASE_ID,
Property.create("eeny", 100L),
Property.create("meeny", 200L),
Property.create("miney", 300L));
ComparableEntity e3 =
builder.addEntityProto(
BASE_ID + 2,
Property.create("moxy", 100L),
Property.create("fruvis", 200L),
Property.create("cortex", 300L));
builder.build();
ImmutableSet<ComparableEntity> entities =
new RecordAccumulator().readDirectory(subdir).getComparableEntitySet();
assertThat(entities).containsExactly(e1, e2, e3);
}
/** Utility class for building a leveldb logfile. */
private static final class LevelDbFileBuilder {
private final FileOutputStream out;
private byte[] currentBlock = new byte[BLOCK_SIZE];
// Write position in the current block.
private int currentPos = 0;
LevelDbFileBuilder(File file) throws FileNotFoundException {
out = new FileOutputStream(file);
}
/**
* Adds a record containing a new entity protobuf to the file.
*
* <p>Returns the ComparableEntity object rather than "this" so that we can check for the
* presence of the entity in the result set.
*/
private ComparableEntity addEntityProto(int id, Property... properties) throws IOException {
Entity entity = new Entity(TEST_ENTITY_KIND, id);
for (Property prop : properties) {
entity.setProperty(prop.name(), prop.value());
}
EntityProto proto = EntityTranslator.convertToPb(entity);
byte[] protoBytes = proto.toByteArray();
if (protoBytes.length > BLOCK_SIZE - currentPos) {
out.write(currentBlock);
currentBlock = new byte[BLOCK_SIZE];
}
currentPos = LevelDbUtil.addRecord(currentBlock, currentPos, ChunkType.FULL, protoBytes);
return new ComparableEntity(entity);
}
/** Writes all remaining data and closes the block. */
void build() throws IOException {
out.write(currentBlock);
out.close();
}
}
@AutoValue
abstract static class Property {
static Property create(String name, Object value) {
return new AutoValue_RecordAccumulatorTest_Property(name, value);
}
abstract String name();
abstract Object value();
}
}