Skip to content

Commit cf76df4

Browse files
suchapalaversuchapalaver
committed
fix(tap-agent): update id types for network version
Signed-off-by: Joseph Livesey <[email protected]>
1 parent 547bb13 commit cf76df4

File tree

3 files changed

+52
-71
lines changed

3 files changed

+52
-71
lines changed

‎crates/tap-agent/src/tap/context/rav.rs‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -386,8 +386,7 @@ mod test {
386386
{
387387
// Insert a rav
388388
let mut new_rav = T::create_rav(
389-
Some(ALLOCATION_ID_0),
390-
None,
389+
ALLOCATION_ID_0,
391390
SIGNER.0.clone(),
392391
TIMESTAMP_NS,
393392
VALUE_AGGREGATE,
@@ -403,8 +402,7 @@ mod test {
403402
// Update the RAV 3 times in quick succession
404403
for i in 0..3 {
405404
new_rav = T::create_rav(
406-
Some(ALLOCATION_ID_0),
407-
None,
405+
ALLOCATION_ID_0,
408406
SIGNER.0.clone(),
409407
TIMESTAMP_NS + i,
410408
VALUE_AGGREGATE - (i as u128),

‎crates/tap-agent/src/tap/context/receipt.rs‎

Lines changed: 28 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -462,17 +462,11 @@ mod test {
462462
#[future(awt)]
463463
context: TapAgentContext<T>,
464464
) where
465-
T: CreateReceipt,
465+
T: CreateReceipt<Id = Address>,
466466
TapAgentContext<T>: ReceiptRead<TapReceipt> + ReceiptDelete,
467467
{
468-
let received_receipt = T::create_received_receipt(
469-
Some(ALLOCATION_ID_0),
470-
None,
471-
&SIGNER.0,
472-
u64::MAX,
473-
u64::MAX,
474-
u128::MAX,
475-
);
468+
let received_receipt =
469+
T::create_received_receipt(ALLOCATION_ID_0, &SIGNER.0, u64::MAX, u64::MAX, u128::MAX);
476470

477471
// Storing the receipt
478472
store_receipt(&context.pgpool, received_receipt.signed_receipt())
@@ -508,9 +502,16 @@ mod test {
508502
let received_receipt_vec: Vec<_> = received_receipt_vec
509503
.iter()
510504
.filter(|received_receipt| {
505+
use thegraph_core::CollectionId;
506+
let expected_collection_id = *CollectionId::from(storage_adapter.allocation_id);
507+
508+
let id_matches = received_receipt.signed_receipt().allocation_id()
509+
== Some(storage_adapter.allocation_id)
510+
|| received_receipt.signed_receipt().collection_id()
511+
== Some(expected_collection_id);
512+
511513
range.contains(&received_receipt.signed_receipt().timestamp_ns())
512-
&& (received_receipt.signed_receipt().allocation_id()
513-
== Some(storage_adapter.allocation_id))
514+
&& id_matches
514515
&& escrow_accounts_snapshot
515516
.get_sender_for_signer(
516517
&received_receipt
@@ -579,13 +580,15 @@ mod test {
579580
.zip(received_receipt_vec.iter())
580581
.collect::<Vec<_>>();
581582

582-
// Remove the received receipts by timestamp range for the correct (allocation_id,
583+
// Remove the received receipts by timestamp range for the correct (collection_id,
583584
// sender)
584585
let received_receipt_vec: Vec<_> = received_receipt_vec
585586
.iter()
586587
.filter(|(_, received_receipt)| {
587-
if (received_receipt.signed_receipt().allocation_id()
588-
== Some(storage_adapter.allocation_id))
588+
use thegraph_core::CollectionId;
589+
let expected_collection_id = *CollectionId::from(storage_adapter.allocation_id);
590+
if (received_receipt.signed_receipt().collection_id()
591+
== Some(expected_collection_id))
589592
&& escrow_accounts_snapshot
590593
.get_sender_for_signer(
591594
&received_receipt
@@ -837,14 +840,13 @@ mod test {
837840
#[future(awt)]
838841
context: TapAgentContext<T>,
839842
) where
840-
T: CreateReceipt,
843+
T: CreateReceipt<Id = Address>,
841844
TapAgentContext<T>: ReceiptRead<TapReceipt> + ReceiptDelete,
842845
{
843846
// Creating 100 receipts with timestamps 42 to 141
844847
for i in 0..100 {
845848
let receipt = T::create_received_receipt(
846-
Some(ALLOCATION_ID_0),
847-
None,
849+
ALLOCATION_ID_0,
848850
&SIGNER.0,
849851
i + 684,
850852
i + 42,
@@ -870,8 +872,7 @@ mod test {
870872
// add a copy in the same timestamp
871873
for i in 0..100 {
872874
let receipt = T::create_received_receipt(
873-
Some(ALLOCATION_ID_0),
874-
None,
875+
ALLOCATION_ID_0,
875876
&SIGNER.0,
876877
i + 684,
877878
i + 43,
@@ -906,15 +907,14 @@ mod test {
906907
#[future(awt)]
907908
context: TapAgentContext<T>,
908909
) where
909-
T: CreateReceipt,
910+
T: CreateReceipt<Id = Address>,
910911
TapAgentContext<T>: ReceiptRead<TapReceipt> + ReceiptDelete,
911912
{
912913
// Creating 10 receipts with timestamps 42 to 51
913914
let mut received_receipt_vec = Vec::new();
914915
for i in 0..10 {
915916
received_receipt_vec.push(T::create_received_receipt(
916-
Some(ALLOCATION_ID_0),
917-
None,
917+
ALLOCATION_ID_0,
918918
&SIGNER.0,
919919
i + 684,
920920
i + 42,
@@ -923,16 +923,14 @@ mod test {
923923

924924
// Adding irrelevant receipts to make sure they are not retrieved
925925
received_receipt_vec.push(T::create_received_receipt(
926-
Some(ALLOCATION_ID_IRRELEVANT),
927-
None,
926+
ALLOCATION_ID_IRRELEVANT,
928927
&SIGNER.0,
929928
i + 684,
930929
i + 42,
931930
(i + 124).into(),
932931
));
933932
received_receipt_vec.push(T::create_received_receipt(
934-
Some(ALLOCATION_ID_0),
935-
None,
933+
ALLOCATION_ID_0,
936934
&SENDER_IRRELEVANT.0,
937935
i + 684,
938936
i + 42,
@@ -1034,15 +1032,14 @@ mod test {
10341032
#[future(awt)]
10351033
context: TapAgentContext<T>,
10361034
) where
1037-
T: CreateReceipt + RemoveRange,
1035+
T: CreateReceipt<Id = Address> + RemoveRange,
10381036
TapAgentContext<T>: ReceiptRead<TapReceipt> + ReceiptDelete,
10391037
{
10401038
// Creating 10 receipts with timestamps 42 to 51
10411039
let mut received_receipt_vec = Vec::new();
10421040
for i in 0..10 {
10431041
received_receipt_vec.push(T::create_received_receipt(
1044-
Some(ALLOCATION_ID_0),
1045-
None,
1042+
ALLOCATION_ID_0,
10461043
&SIGNER.0,
10471044
i + 684,
10481045
i + 42,
@@ -1051,16 +1048,14 @@ mod test {
10511048

10521049
// Adding irrelevant receipts to make sure they are not retrieved
10531050
received_receipt_vec.push(T::create_received_receipt(
1054-
Some(ALLOCATION_ID_IRRELEVANT),
1055-
None,
1051+
ALLOCATION_ID_IRRELEVANT,
10561052
&SIGNER.0,
10571053
i + 684,
10581054
i + 42,
10591055
(i + 124).into(),
10601056
));
10611057
received_receipt_vec.push(T::create_received_receipt(
1062-
Some(ALLOCATION_ID_0),
1063-
None,
1058+
ALLOCATION_ID_0,
10641059
&SENDER_IRRELEVANT.0,
10651060
i + 684,
10661061
i + 42,

‎crates/tap-agent/src/test.rs‎

Lines changed: 22 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -262,14 +262,10 @@ pub async fn create_sender_accounts_manager(
262262
)
263263
}
264264

265-
/// Generic implementation of create_rav
265+
/// Network-version specific RAV creation
266266
pub trait CreateRav: NetworkVersion {
267-
/// This might seem weird at first glance since [Horizon] and [Legacy] implementation have the same
268-
/// function signature and don't require &self. The reason is that we can not match over T to get
269-
/// all variants because T is a trait and not an enum.
270267
fn create_rav(
271-
allocation_id: Option<Address>,
272-
collection_id: Option<FixedBytes<32>>,
268+
id: Address,
273269
signer_wallet: PrivateKeySigner,
274270
timestamp_ns: u64,
275271
value_aggregate: u128,
@@ -278,35 +274,25 @@ pub trait CreateRav: NetworkVersion {
278274

279275
impl CreateRav for Legacy {
280276
fn create_rav(
281-
allocation_id: Option<Address>,
282-
_collection_id: Option<FixedBytes<32>>,
277+
allocation_id: Address,
283278
signer_wallet: PrivateKeySigner,
284279
timestamp_ns: u64,
285280
value_aggregate: u128,
286281
) -> Eip712SignedMessage<Self::Rav> {
287-
create_rav(
288-
allocation_id.unwrap(),
289-
signer_wallet,
290-
timestamp_ns,
291-
value_aggregate,
292-
)
282+
create_rav(allocation_id, signer_wallet, timestamp_ns, value_aggregate)
293283
}
294284
}
295285

296286
impl CreateRav for Horizon {
297287
fn create_rav(
298-
_allocation_id: Option<Address>,
299-
collection_id: Option<FixedBytes<32>>,
288+
allocation_id: Address,
300289
signer_wallet: PrivateKeySigner,
301290
timestamp_ns: u64,
302291
value_aggregate: u128,
303292
) -> Eip712SignedMessage<Self::Rav> {
304-
create_rav_v2(
305-
collection_id.unwrap(),
306-
signer_wallet,
307-
timestamp_ns,
308-
value_aggregate,
309-
)
293+
use thegraph_core::CollectionId;
294+
let collection_id = *CollectionId::from(allocation_id);
295+
create_rav_v2(collection_id, signer_wallet, timestamp_ns, value_aggregate)
310296
}
311297
}
312298

@@ -352,14 +338,12 @@ pub fn create_rav_v2(
352338
.unwrap()
353339
}
354340

355-
/// Generic implementation of create_received_receipt
341+
/// Network-version specific receipt creation
356342
pub trait CreateReceipt {
357-
/// This might seem weird at first glance since [Horizon] and [Legacy] implementation have the same
358-
/// function signature and don't require &self. The reason is that we can not match over T to get
359-
/// all variants because T is a trait and not an enum.
343+
type Id: Clone + std::fmt::Debug;
344+
360345
fn create_received_receipt(
361-
allocation_id: Option<Address>,
362-
collection_id: Option<FixedBytes<32>>,
346+
id: Self::Id,
363347
signer_wallet: &PrivateKeySigner,
364348
nonce: u64,
365349
timestamp_ns: u64,
@@ -368,18 +352,21 @@ pub trait CreateReceipt {
368352
}
369353

370354
impl CreateReceipt for Horizon {
355+
type Id = Address;
356+
371357
fn create_received_receipt(
372-
_allocation_id: Option<Address>,
373-
collection_id: Option<FixedBytes<32>>,
358+
allocation_id: Self::Id,
374359
signer_wallet: &PrivateKeySigner,
375360
nonce: u64,
376361
timestamp_ns: u64,
377362
value: u128,
378363
) -> CheckingReceipt {
364+
use thegraph_core::CollectionId;
365+
let collection_id = *CollectionId::from(allocation_id);
379366
let receipt = Eip712SignedMessage::new(
380367
&TAP_EIP712_DOMAIN_SEPARATOR,
381368
tap_graph::v2::Receipt {
382-
collection_id: collection_id.unwrap(),
369+
collection_id,
383370
payer: SENDER.1,
384371
service_provider: INDEXER.1,
385372
data_service: Address::ZERO,
@@ -395,9 +382,10 @@ impl CreateReceipt for Horizon {
395382
}
396383

397384
impl CreateReceipt for Legacy {
385+
type Id = Address;
386+
398387
fn create_received_receipt(
399-
allocation_id: Option<Address>,
400-
_collection_id: Option<FixedBytes<32>>,
388+
allocation_id: Self::Id,
401389
signer_wallet: &PrivateKeySigner,
402390
nonce: u64,
403391
timestamp_ns: u64,
@@ -406,7 +394,7 @@ impl CreateReceipt for Legacy {
406394
let receipt = Eip712SignedMessage::new(
407395
&TAP_EIP712_DOMAIN_SEPARATOR,
408396
Receipt {
409-
allocation_id: allocation_id.unwrap(),
397+
allocation_id,
410398
nonce,
411399
timestamp_ns,
412400
value,

0 commit comments

Comments
 (0)