tests: udpate subgraph composition integration test cases to handle entity ops

Author: incrypto32

graph/src/runtime/mod.rs

@@ -368,11 +368,9 @@ pub enum IndexForAscTypeId {
     // ...
     // LastStarknetType = 4499,
 
-
     // Subgraph Data Source types
     AscEntityTrigger = 4500,
 
-
     // Reserved discriminant space for a future blockchain type IDs: [4,500, 5,499]
     //
     // Generated with the following shell script:

tests/integration-tests/source-subgraph/schema.graphql

@@ -1,12 +1,13 @@
-
 type Block @entity {
   id: ID!
   number: BigInt!
   hash: Bytes!
+  testMessage: String
 }
 
 type Block2 @entity {
   id: ID!
   number: BigInt!
   hash: Bytes!
+  testMessage: String
 }

tests/integration-tests/source-subgraph/src/mapping.ts

@@ -1,4 +1,4 @@
-import { ethereum, log } from '@graphprotocol/graph-ts';
+import { ethereum, log, store } from '@graphprotocol/graph-ts';
 import { Block, Block2 } from '../generated/schema';
 import { BigInt } from '@graphprotocol/graph-ts';
 
@@ -22,4 +22,36 @@ export function handleBlock(block: ethereum.Block): void {
   blockEntity3.number = block.number;
   blockEntity3.hash = block.hash;
   blockEntity3.save();
+
+  if (block.number.equals(BigInt.fromI32(1))) {
+    let id = 'TEST';
+    let entity = new Block(id);
+    entity.number = block.number;
+    entity.hash = block.hash;
+    entity.testMessage = 'Created at block 1';
+    log.info('Created entity at block 1', []);
+    entity.save();
+  }
+
+  if (block.number.equals(BigInt.fromI32(2))) {
+    let id = 'TEST';
+    let blockEntity1 = Block.load(id);
+    if (blockEntity1) {
+      // Update the block entity
+      blockEntity1.testMessage = 'Updated at block 2';
+      log.info('Updated entity at block 2', []);
+      blockEntity1.save();
+    }
+  }
+
+  if (block.number.equals(BigInt.fromI32(3))) {
+    let id = 'TEST';
+    let blockEntity1 = Block.load(id);
+    if (blockEntity1) {
+      blockEntity1.testMessage = 'Deleted at block 3';
+      log.info('Deleted entity at block 3', []);
+      blockEntity1.save();
+      store.remove('Block', id);
+    }
+  }
 }

tests/integration-tests/subgraph-data-sources/schema.graphql

@@ -2,4 +2,5 @@ type MirrorBlock @entity {
   id: String!
   number: BigInt!
   hash: Bytes!
+  testMessage: String
 }

tests/integration-tests/subgraph-data-sources/src/mapping.ts

@@ -1,4 +1,4 @@
-import { Entity, log } from '@graphprotocol/graph-ts';
+import { Entity, log, store } from '@graphprotocol/graph-ts';
 import { MirrorBlock } from '../generated/schema';
 
 export class EntityTrigger {
@@ -14,12 +14,33 @@ export function handleEntity(trigger: EntityTrigger): void {
   let blockEntity = trigger.entity;
   let blockNumber = blockEntity.getBigInt('number');
   let blockHash = blockEntity.getBytes('hash');
+  let testMessage = blockEntity.get('testMessage');
   let id = blockEntity.getString('id');
 
   log.info('Block number: {}', [blockNumber.toString()]);
 
-  let block = new MirrorBlock(id);
+  if (trigger.entityOp == 2) {
+    log.info('Removing block entity with id: {}', [id]);
+    store.remove('MirrorBlock', id);
+    return;
+  }
+
+  let block = loadOrCreateMirrorBlock(id);
   block.number = blockNumber;
   block.hash = blockHash;
+  if (testMessage) {
+    block.testMessage = testMessage.toString();
+  }
+
   block.save();
 }
+
+export function loadOrCreateMirrorBlock(id: string): MirrorBlock {
+  let block = MirrorBlock.load(id);
+  if (!block) {
+    log.info('Creating new block entity with id: {}', [id]);
+    block = new MirrorBlock(id);
+  }
+
+  return block;
+}

tests/integration-tests/subgraph-data-sources/subgraph.yaml

@@ -6,7 +6,7 @@ dataSources:
     name: Contract
     network: test
     source:
-      address: 'QmeZhEiJuBusu7GxCe6AytvqSsgwV8QxkbSYx5ojSFB28a'
+      address: 'Qmaqf8cRxfxbduZppSHKG9DMuX5JZPMoGuwGb2DQuo48sq'
       startBlock: 0
     mapping:
       apiVersion: 0.0.7

tests/tests/integration_tests.rs

@@ -563,6 +563,42 @@ async fn subgraph_data_sources(ctx: TestContext) -> anyhow::Result<()> {
     )
     .await?;
 
+    let expected_response = json!({
+        "mirrorBlock": { "id": "TEST", "number": "1", "testMessage": "Created at block 1" },
+    });
+
+    query_succeeds(
+        "Blocks should be right",
+        &subgraph,
+        "{ mirrorBlock(id: \"TEST\", block: {number: 1}) { id, number, testMessage } }",
+        expected_response,
+    )
+    .await?;
+
+    let expected_response = json!({
+        "mirrorBlock": { "id": "TEST", "number": "1", "testMessage": "Updated at block 2" },
+    });
+
+    query_succeeds(
+        "Blocks should be right",
+        &subgraph,
+        "{ mirrorBlock(id: \"TEST\", block: {number: 2}) { id, number, testMessage } }",
+        expected_response,
+    )
+    .await?;
+
+    let expected_response = json!({
+        "mirrorBlock": null,
+    });
+
+    query_succeeds(
+        "Blocks should be right",
+        &subgraph,
+        "{ mirrorBlock(id: \"TEST\", block: {number: 3}) { id, number, testMessage } }",
+        expected_response,
+    )
+    .await?;
+
     Ok(())
 }