fix: bug where newtotalstake was not correct, address feedback

Author: kemuru

subgraph/core-university/src/SortitionModule.ts

@@ -1,26 +1,25 @@
 import { SortitionModule, StakeLocked, StakeSet as StakeSetEvent } from "../generated/SortitionModule/SortitionModule";
-import { StakeSet as StakeSetEntity, User } from "../generated/schema";
+import { StakeSet as StakeSetEntity } from "../generated/schema";
 
-import { updateJurorDelayedStake, updateJurorStake } from "./entities/JurorTokensPerCourt";
+import { ensureJurorTokensPerCourt, updateJurorDelayedStake, updateJurorStake } from "./entities/JurorTokensPerCourt";
 import { ensureUser } from "./entities/User";
 import { ZERO } from "./utils";
 
 export function handleStakeSet(event: StakeSetEvent): void {
   const jurorAddress = event.params._address.toHexString();
-  ensureUser(jurorAddress);
+  const juror = ensureUser(jurorAddress);
   const courtID = event.params._courtID.toString();
 
   updateJurorStake(jurorAddress, courtID.toString(), SortitionModule.bind(event.address), event.block.timestamp);
   //stake is updated instantly so no delayed amount, set delay amount to zero
   updateJurorDelayedStake(jurorAddress, courtID, ZERO);
 
-  const juror = User.load(jurorAddress);
-  if (!juror) return;
   const stakeSet = new StakeSetEntity(event.transaction.hash.toHex() + "-" + event.logIndex.toString());
-  stakeSet.address = jurorAddress;
+  const jurorTokensPerCourt = ensureJurorTokensPerCourt(jurorAddress, "1");
+  stakeSet.juror = juror.id;
   stakeSet.courtID = event.params._courtID;
   stakeSet.stake = event.params._amount;
-  stakeSet.newTotalStake = juror.totalStake;
+  stakeSet.newTotalStake = jurorTokensPerCourt.effectiveStake;
   stakeSet.blocknumber = event.block.number;
   stakeSet.timestamp = event.block.timestamp;
   stakeSet.logIndex = event.logIndex;

subgraph/core/schema.graphql

@@ -354,9 +354,9 @@ type ClassicContribution implements Contribution @entity {
   rewardWithdrawn: Boolean!
 }
 
-type StakeSet @entity {
+type StakeSet @entity(immutable: true) {
   id: ID! # event.transaction.hash.toHex() + - + event.logIndex.toString()
-  address: String!
+  juror: User!
   courtID: BigInt!
   stake: BigInt!
   newTotalStake: BigInt!

subgraph/core/src/SortitionModule.ts

@@ -6,9 +6,9 @@ import {
   StakeLocked,
   StakeSet as StakeSetEvent,
 } from "../generated/SortitionModule/SortitionModule";
-import { StakeSet as StakeSetEntity, User } from "../generated/schema";
+import { StakeSet as StakeSetEntity } from "../generated/schema";
 
-import { updateJurorDelayedStake, updateJurorStake } from "./entities/JurorTokensPerCourt";
+import { ensureJurorTokensPerCourt, updateJurorDelayedStake, updateJurorStake } from "./entities/JurorTokensPerCourt";
 import { ensureUser } from "./entities/User";
 import { ZERO } from "./utils";
 
@@ -26,20 +26,19 @@ export function handleStakeDelayedNotTransferred(event: StakeDelayedNotTransferr
 
 export function handleStakeSet(event: StakeSetEvent): void {
   const jurorAddress = event.params._address.toHexString();
-  ensureUser(jurorAddress);
+  const juror = ensureUser(jurorAddress);
   const courtID = event.params._courtID.toString();
 
   updateJurorStake(jurorAddress, courtID.toString(), SortitionModule.bind(event.address), event.block.timestamp);
   //stake is updated instantly so no delayed amount, set delay amount to zero
   updateJurorDelayedStake(jurorAddress, courtID, ZERO);
 
-  const juror = User.load(jurorAddress);
-  if (!juror) return;
   const stakeSet = new StakeSetEntity(event.transaction.hash.toHex() + "-" + event.logIndex.toString());
-  stakeSet.address = jurorAddress;
+  const jurorTokensPerCourt = ensureJurorTokensPerCourt(jurorAddress, "1");
+  stakeSet.juror = juror.id;
   stakeSet.courtID = event.params._courtID;
   stakeSet.stake = event.params._amount;
-  stakeSet.newTotalStake = juror.totalStake;
+  stakeSet.newTotalStake = jurorTokensPerCourt.effectiveStake;
   stakeSet.blocknumber = event.block.number;
   stakeSet.timestamp = event.block.timestamp;
   stakeSet.logIndex = event.logIndex;