Skip to content

Commit 07e8f3b

Browse files
committed
fix(tap-agent): update id types for network version
1 parent 547bb13 commit 07e8f3b

File tree

3 files changed

+43
-69
lines changed

3 files changed

+43
-69
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: 19 additions & 31 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())
@@ -579,13 +573,15 @@ mod test {
579573
.zip(received_receipt_vec.iter())
580574
.collect::<Vec<_>>();
581575

582-
// Remove the received receipts by timestamp range for the correct (allocation_id,
576+
// Remove the received receipts by timestamp range for the correct (collection_id,
583577
// sender)
584578
let received_receipt_vec: Vec<_> = received_receipt_vec
585579
.iter()
586580
.filter(|(_, received_receipt)| {
587-
if (received_receipt.signed_receipt().allocation_id()
588-
== Some(storage_adapter.allocation_id))
581+
use thegraph_core::CollectionId;
582+
let expected_collection_id = *CollectionId::from(storage_adapter.allocation_id);
583+
if (received_receipt.signed_receipt().collection_id()
584+
== Some(expected_collection_id))
589585
&& escrow_accounts_snapshot
590586
.get_sender_for_signer(
591587
&received_receipt
@@ -837,14 +833,13 @@ mod test {
837833
#[future(awt)]
838834
context: TapAgentContext<T>,
839835
) where
840-
T: CreateReceipt,
836+
T: CreateReceipt<Id = Address>,
841837
TapAgentContext<T>: ReceiptRead<TapReceipt> + ReceiptDelete,
842838
{
843839
// Creating 100 receipts with timestamps 42 to 141
844840
for i in 0..100 {
845841
let receipt = T::create_received_receipt(
846-
Some(ALLOCATION_ID_0),
847-
None,
842+
ALLOCATION_ID_0,
848843
&SIGNER.0,
849844
i + 684,
850845
i + 42,
@@ -870,8 +865,7 @@ mod test {
870865
// add a copy in the same timestamp
871866
for i in 0..100 {
872867
let receipt = T::create_received_receipt(
873-
Some(ALLOCATION_ID_0),
874-
None,
868+
ALLOCATION_ID_0,
875869
&SIGNER.0,
876870
i + 684,
877871
i + 43,
@@ -906,15 +900,14 @@ mod test {
906900
#[future(awt)]
907901
context: TapAgentContext<T>,
908902
) where
909-
T: CreateReceipt,
903+
T: CreateReceipt<Id = Address>,
910904
TapAgentContext<T>: ReceiptRead<TapReceipt> + ReceiptDelete,
911905
{
912906
// Creating 10 receipts with timestamps 42 to 51
913907
let mut received_receipt_vec = Vec::new();
914908
for i in 0..10 {
915909
received_receipt_vec.push(T::create_received_receipt(
916-
Some(ALLOCATION_ID_0),
917-
None,
910+
ALLOCATION_ID_0,
918911
&SIGNER.0,
919912
i + 684,
920913
i + 42,
@@ -923,16 +916,14 @@ mod test {
923916

924917
// Adding irrelevant receipts to make sure they are not retrieved
925918
received_receipt_vec.push(T::create_received_receipt(
926-
Some(ALLOCATION_ID_IRRELEVANT),
927-
None,
919+
ALLOCATION_ID_IRRELEVANT,
928920
&SIGNER.0,
929921
i + 684,
930922
i + 42,
931923
(i + 124).into(),
932924
));
933925
received_receipt_vec.push(T::create_received_receipt(
934-
Some(ALLOCATION_ID_0),
935-
None,
926+
ALLOCATION_ID_0,
936927
&SENDER_IRRELEVANT.0,
937928
i + 684,
938929
i + 42,
@@ -1034,15 +1025,14 @@ mod test {
10341025
#[future(awt)]
10351026
context: TapAgentContext<T>,
10361027
) where
1037-
T: CreateReceipt + RemoveRange,
1028+
T: CreateReceipt<Id = Address> + RemoveRange,
10381029
TapAgentContext<T>: ReceiptRead<TapReceipt> + ReceiptDelete,
10391030
{
10401031
// Creating 10 receipts with timestamps 42 to 51
10411032
let mut received_receipt_vec = Vec::new();
10421033
for i in 0..10 {
10431034
received_receipt_vec.push(T::create_received_receipt(
1044-
Some(ALLOCATION_ID_0),
1045-
None,
1035+
ALLOCATION_ID_0,
10461036
&SIGNER.0,
10471037
i + 684,
10481038
i + 42,
@@ -1051,16 +1041,14 @@ mod test {
10511041

10521042
// Adding irrelevant receipts to make sure they are not retrieved
10531043
received_receipt_vec.push(T::create_received_receipt(
1054-
Some(ALLOCATION_ID_IRRELEVANT),
1055-
None,
1044+
ALLOCATION_ID_IRRELEVANT,
10561045
&SIGNER.0,
10571046
i + 684,
10581047
i + 42,
10591048
(i + 124).into(),
10601049
));
10611050
received_receipt_vec.push(T::create_received_receipt(
1062-
Some(ALLOCATION_ID_0),
1063-
None,
1051+
ALLOCATION_ID_0,
10641052
&SENDER_IRRELEVANT.0,
10651053
i + 684,
10661054
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)