address CI failures

Author: MoonBoi9001

runtime/wasm/src/host_exports.rs

@@ -33,6 +33,8 @@ use crate::{error::DeterminismLevel, module::IntoTrap};
 
 use super::module::WasmInstanceData;
 
+use graph::schema::EntityKey;
+
 fn write_poi_event(
     proof_of_indexing: &SharedProofOfIndexing,
     poi_event: &ProofOfIndexingEvent,
@@ -354,7 +356,7 @@ impl HostExports {
             .metrics
             .track_storage_size_change(&entity_type, &entity, false);
 
-        if !state.entity_cache.contains_key(&key) {
+        if state.entity_cache.get(&key, GetScope::Store)?.is_none() {
             state.metrics.track_entity_count_change(&entity_type, 1);
         }
 
@@ -1250,14 +1252,15 @@ impl HostExports {
             .context("Failed to decode")
     }
 
-    fn remove_entity(&mut self, key: &EntityKey) -> Result<(), HostExportError> {
+    fn remove_entity(
+        &mut self,
+        key: &EntityKey,
+        state: &mut BlockState,
+    ) -> Result<(), HostExportError> {
         let entity_type = key.entity_type.clone();
 
-        if let Some(entity) = state.entity_cache.get(&key, GetScope::Store)? {
-            // Track removal metrics while we still have the entity data
-            state
-                .metrics
-                .track_storage_size_change(&entity_type, &entity, true);
+        if let Some(entity) = state.entity_cache.get(key, GetScope::Store)? {
+            state.metrics.track_storage_size_change(&entity_type, &entity, true);
             state.metrics.track_entity_count_change(&entity_type, -1);
         }
 
@@ -1395,3 +1398,9 @@ fn bytes_to_string_is_lossy() {
         )
     )
 }
+
+impl From<StoreError> for HostExportError {
+    fn from(err: StoreError) -> Self {
+        HostExportError::Deterministic(DeterministicHostError::from(err))
+    }
+}