Add tests and fix middleware

crates/client/src/client.rs

@@ -240,6 +240,28 @@ impl FuelClient {
         Self::decode_response(response)
     }
 
+    pub async fn query_with_headers<ResponseData, Vars>(
+        &self,
+        q: Operation<ResponseData, Vars>,
+        headers: impl IntoIterator<Item = (String, String)>,
+    ) -> io::Result<ResponseData>
+    where
+        Vars: serde::Serialize,
+        ResponseData: serde::de::DeserializeOwned + 'static,
+    {
+        let mut request_builder = self.client.post(self.url.clone());
+        for (header_name, header_value) in headers {
+            request_builder = request_builder.header(header_name, header_value);
+        }
+
+        let response = request_builder
+            .run_graphql(q)
+            .await
+            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+
+        Self::decode_response(response)
+    }
+
     fn decode_response<R>(response: GraphQlResponse<R>) -> io::Result<R>
     where
         R: serde::de::DeserializeOwned + 'static,
@@ -1080,6 +1102,32 @@ impl FuelClient {
         Ok(balance.amount)
     }
 
+    pub async fn balance_with_required_block_header(
+        &self,
+        owner: &Address,
+        asset_id: Option<&AssetId>,
+        required_block_height: u32,
+    ) -> io::Result<u128> {
+        let owner: schema::Address = (*owner).into();
+        let asset_id: schema::AssetId = match asset_id {
+            Some(asset_id) => (*asset_id).into(),
+            None => schema::AssetId::default(),
+        };
+        let query = schema::balance::BalanceQuery::build(BalanceArgs { owner, asset_id });
+        let balance: types::Balance = self
+            .query_with_headers(
+                query,
+                vec![(
+                    "REQUIRED_FUEL_BLOCK_HEIGHT".to_string(),
+                    required_block_height.to_string(),
+                )],
+            )
+            .await?
+            .balance
+            .into();
+        Ok(balance.amount)
+    }
+
     // Retrieve a page of balances by their owner
     pub async fn balances(
         &self,

crates/fuel-core/src/graphql_api/api_service.rs

@@ -258,11 +258,12 @@ async fn required_fuel_block_height<B>(
     let raw_required_fuel_block_height =
         required_fuel_block_height_header.map_err(|_err| StatusCode::BAD_REQUEST)?;
 
-    let required_fuel_block_height = raw_required_fuel_block_height
-        .parse::<BlockHeight>()
-        .map_err(|_| StatusCode::BAD_REQUEST)?;
+    let required_fuel_block_height: BlockHeight = raw_required_fuel_block_height
+        .parse::<u32>()
+        .map_err(|_| StatusCode::BAD_REQUEST)?
+        .into();
 
-    if required_fuel_block_height < last_known_block_height {
+    if required_fuel_block_height > last_known_block_height {
         Err(StatusCode::PRECONDITION_FAILED)
     } else {
         let mut response = next.run(req).await;

tests/tests/lib.rs

@@ -48,6 +48,8 @@ mod regenesis;
 #[cfg(not(feature = "only-p2p"))]
 mod relayer;
 #[cfg(not(feature = "only-p2p"))]
+mod required_fuel_block_height_header;
+#[cfg(not(feature = "only-p2p"))]
 mod snapshot;
 #[cfg(not(feature = "only-p2p"))]
 mod state_rewind;

tests/tests/required_fuel_block_height_header.rs

@@ -0,0 +1,50 @@
+use fuel_core::{
+    chain_config::StateConfig,
+    service::{
+        Config,
+        FuelService,
+    },
+};
+use fuel_core_client::client::{
+    types::primitives::{
+        Address,
+        AssetId,
+    },
+    FuelClient,
+};
+
+#[tokio::test]
+async fn balance_with_block_height_header() {
+    let owner = Address::default();
+    let asset_id = AssetId::BASE;
+
+    // setup config
+    let state_config = StateConfig::default();
+    let config = Config::local_node_with_state_config(state_config);
+
+    // setup server & client
+    let srv = FuelService::new_node(config).await.unwrap();
+    let client = FuelClient::from(srv.bound_address);
+
+    // Issue a request with wrong precondition
+    let error = client
+        .balance_with_required_block_header(&owner, Some(&asset_id), 100)
+        .await
+        .unwrap_err();
+
+    let error_str = format!("{:?}", error);
+    assert_eq!(
+        error_str,
+        "Custom { kind: Other, error: ErrorResponse(412, \"\") }"
+    );
+
+    // Meet precondition on server side
+    client.produce_blocks(100, None).await.unwrap();
+
+    // Issue request again
+    let result = client
+        .balance_with_required_block_header(&owner, Some(&asset_id), 100)
+        .await;
+
+    assert!(result.is_ok());
+}