@@ -11,22 +11,18 @@ import (
1111)
1212
1313func TestUsageStore_All (t * testing.T ) {
14- now := time .Now ()
15- cutoff := now .Add (- 5 * time .Minute ).UnixNano ()
1614 // Create a store with 10 partitions.
17- s := NewUsageStore (Config {NumPartitions : 10 })
15+ s := NewUsageStore (Config {
16+ NumPartitions : 10 ,
17+ WindowSize : time .Minute ,
18+ })
19+ clock := quartz .NewMock (t )
20+ s .clock = clock
1821 // Create 10 streams. Since we use i as the hash, we can expect the
1922 // streams to be sharded over all 10 partitions.
20- streams := make ([]* proto.StreamMetadata , 10 )
2123 for i := 0 ; i < 10 ; i ++ {
22- streams [i ] = & proto.StreamMetadata {
23- StreamHash : uint64 (i ),
24- }
24+ s .set ("tenant" , Stream {Hash : uint64 (i )})
2525 }
26- // Add the streams to the store, all streams should be accepted.
27- accepted , rejected := s .Update ("tenant" , streams , now .UnixNano (), cutoff , 0 , 0 , nil )
28- require .Len (t , accepted , 10 )
29- require .Empty (t , rejected )
3026 // Check that we can iterate all stored streams.
3127 expected := []uint64 {0x0 , 0x1 , 0x2 , 0x3 , 0x4 , 0x5 , 0x6 , 0x7 , 0x8 , 0x9 }
3228 actual := make ([]uint64 , 0 , len (expected ))
@@ -37,26 +33,22 @@ func TestUsageStore_All(t *testing.T) {
3733}
3834
3935func TestUsageStore_ForTenant (t * testing.T ) {
40- now := time .Now ()
41- cutoff := now .Add (- 5 * time .Minute ).UnixNano ()
4236 // Create a store with 10 partitions.
43- s := NewUsageStore (Config {NumPartitions : 10 })
37+ s := NewUsageStore (Config {
38+ NumPartitions : 10 ,
39+ WindowSize : time .Minute ,
40+ })
41+ clock := quartz .NewMock (t )
42+ s .clock = clock
4443 // Create 10 streams. Since we use i as the hash, we can expect the
4544 // streams to be sharded over all 10 partitions.
46- streams := make ([]* proto.StreamMetadata , 10 )
4745 for i := 0 ; i < 10 ; i ++ {
48- streams [i ] = & proto.StreamMetadata {
49- StreamHash : uint64 (i ),
46+ tenant := "tenant1"
47+ if i >= 5 {
48+ tenant = "tenant2"
5049 }
50+ s .set (tenant , Stream {Hash : uint64 (i )})
5151 }
52- // Add the streams to the store, but with the streams shared between
53- // two tenants.
54- accepted , rejected := s .Update ("tenant1" , streams [0 :5 ], now .UnixNano (), cutoff , 0 , 0 , nil )
55- require .Len (t , accepted , 5 )
56- require .Empty (t , rejected )
57- accepted , rejected = s .Update ("tenant2" , streams [5 :], now .UnixNano (), cutoff , 0 , 0 , nil )
58- require .Len (t , accepted , 5 )
59- require .Empty (t , rejected )
6052 // Check we can iterate just the streams for each tenant.
6153 expected1 := []uint64 {0x0 , 0x1 , 0x2 , 0x3 , 0x4 }
6254 actual1 := make ([]uint64 , 0 , 5 )
@@ -73,11 +65,6 @@ func TestUsageStore_ForTenant(t *testing.T) {
7365}
7466
7567func TestUsageStore_Store (t * testing.T ) {
76- now := time .Now ()
77- cutoff := now .Add (- 5 * time .Minute ).UnixNano ()
78- bucketStart := now .Truncate (time .Minute ).UnixNano ()
79- bucketCutoff := now .Add (- 5 * time .Minute ).UnixNano ()
80-
8168 tests := []struct {
8269 name string
8370 numPartitions int
@@ -180,10 +167,15 @@ func TestUsageStore_Store(t *testing.T) {
180167
181168 for _ , test := range tests {
182169 t .Run (test .name , func (t * testing.T ) {
183- s := NewUsageStore (Config {NumPartitions : test .numPartitions })
184- s .Update ("tenant" , test .seed , now .UnixNano (), cutoff , bucketStart , bucketCutoff , nil )
170+ s := NewUsageStore (Config {
171+ NumPartitions : test .numPartitions ,
172+ WindowSize : time .Minute ,
173+ })
174+ clock := quartz .NewMock (t )
175+ s .clock = clock
176+ s .Update ("tenant" , test .seed , clock .Now (), nil )
185177 streamLimitCond := streamLimitExceeded (test .maxGlobalStreams )
186- accepted , rejected := s .Update ("tenant" , test .streams , now . UnixNano (), cutoff , bucketStart , bucketCutoff , streamLimitCond )
178+ accepted , rejected := s .Update ("tenant" , test .streams , clock . Now () , streamLimitCond )
187179 require .ElementsMatch (t , test .expectedAccepted , accepted )
188180 require .ElementsMatch (t , test .expectedRejected , rejected )
189181 })
@@ -225,17 +217,17 @@ func TestUsageStore_Evict(t *testing.T) {
225217
226218func TestUsageStore_EvictPartitions (t * testing.T ) {
227219 // Create a store with 10 partitions.
228- s := NewUsageStore (Config {NumPartitions : 10 })
220+ s := NewUsageStore (Config {
221+ NumPartitions : 10 ,
222+ WindowSize : time .Minute ,
223+ })
224+ clock := quartz .NewMock (t )
225+ s .clock = clock
229226 // Create 10 streams. Since we use i as the hash, we can expect the
230227 // streams to be sharded over all 10 partitions.
231- streams := make ([]* proto.StreamMetadata , 10 )
232228 for i := 0 ; i < 10 ; i ++ {
233- streams [i ] = & proto.StreamMetadata {
234- StreamHash : uint64 (i ),
235- }
229+ s .set ("tenant" , Stream {Hash : uint64 (i )})
236230 }
237- now := time .Now ()
238- s .Update ("tenant" , streams , now .UnixNano (), now .Add (- time .Minute ).UnixNano (), 0 , 0 , nil )
239231 // Evict the first 5 partitions.
240232 s .EvictPartitions ([]int32 {0 , 1 , 2 , 3 , 4 })
241233 // The last 5 partitions should still have data.
0 commit comments