diff --git a/.gitmodules b/.gitmodules
index e684cee..771ae4c 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -1,6 +1,3 @@
-[submodule "lib/forge-std"]
-	path = lib/forge-std
-	url = https://github.com/foundry-rs/forge-std
-[submodule "lib/axiom-v2-periphery"]
-	path = lib/axiom-v2-periphery
-	url = https://github.com/axiom-crypto/axiom-v2-periphery
+[submodule "lib/axiom-std"]
+	path = lib/axiom-std
+	url = https://github.com/axiom-crypto/axiom-std
diff --git a/app/axiom/average.circuit.ts b/app/axiom/average.circuit.ts
index f438ea1..b4f09c6 100644
--- a/app/axiom/average.circuit.ts
+++ b/app/axiom/average.circuit.ts
@@ -12,20 +12,27 @@ import {
   getAccount,
 } from "@axiom-crypto/client";
 
-/// For type safety, define the input types to your circuit here.
-/// These should be the _variable_ inputs to your circuit. Constants can be hard-coded into the circuit itself.
+// For type safety, define the input types to your circuit here.
+// These should be the _variable_ inputs to your circuit. Constants can be hard-coded into the circuit itself.
 export interface CircuitInputs {
   blockNumber: CircuitValue;
   address: CircuitValue;
 }
 
+// Default inputs to use for compiling the circuit. These values should be different than the inputs fed into
+// the circuit at proving time.
+export const defaultInputs = {
+  "blockNumber": 4000000,
+  "address": "0xEaa455e4291742eC362Bc21a8C46E5F2b5ed4701"
+};
+
 // The function name `circuit` is searched for by default by our Axiom CLI; if you decide to 
 // change the function name, you'll also need to ensure that you also pass the Axiom CLI flag 
 // `-f <circuitFunctionName>` for it to work
 export const circuit = async (inputs: CircuitInputs) => {
   // Number of samples to take. Note that this must be a constant value and NOT an input because the size of 
   // the circuit must be known at compile time.
-  const samples = 8; 
+  const samples = 8;
 
   // Number of blocks between each sample.
   const spacing = 900;
diff --git a/app/index.ts b/app/index.ts
index e3f1810..63854f0 100644
--- a/app/index.ts
+++ b/app/index.ts
@@ -15,10 +15,6 @@ const axiomMain = async (input: UserInput<CircuitInputs>) => {
   const axiom = new Axiom({
     circuit: circuit,
     compiledCircuit: compiledCircuit,
-    inputSchema: {
-      blockNumber: "uint32",
-      address: "address",
-    },
     chainId: "11155111",  // Sepolia
     provider: process.env.PROVIDER_URI_SEPOLIA as string,
     privateKey: process.env.PRIVATE_KEY_SEPOLIA as string,
diff --git a/foundry.toml b/foundry.toml
index aea10c1..a5b4684 100644
--- a/foundry.toml
+++ b/foundry.toml
@@ -3,9 +3,8 @@ src = "src"
 out = "out"
 solc = "0.8.19"
 libs = ["node_modules", "lib"]
-remappings = ["@axiom-crypto/v2-periphery/=lib/axiom-v2-periphery/src/", "ds-test/=lib/forge-std/lib/ds-test/src/", "forge-std/=lib/forge-std/src/"]
+remappings = ["@axiom-crypto/axiom-std/=lib/axiom-std/src/"]
 ffi = true
-fs_permissions = [{ access = "read-write", path = ".axiom" }]
 
 # See more config options https://github.com/foundry-rs/foundry/tree/master/config
 [rpc_endpoints]
diff --git a/lib/axiom-std b/lib/axiom-std
new file mode 160000
index 0000000..28faed4
--- /dev/null
+++ b/lib/axiom-std
@@ -0,0 +1 @@
+Subproject commit 28faed462f5f1f21f18765c168bc070671c1a0fe
diff --git a/lib/axiom-v2-periphery b/lib/axiom-v2-periphery
deleted file mode 160000
index 8ce739c..0000000
--- a/lib/axiom-v2-periphery
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit 8ce739ca5a457f85695cede988feb86f1ec703f2
diff --git a/lib/forge-std b/lib/forge-std
deleted file mode 160000
index ae570fe..0000000
--- a/lib/forge-std
+++ /dev/null
@@ -1 +0,0 @@
-Subproject commit ae570fec082bfe1c1f45b0acca4a2b4f84d345ce
diff --git a/package.json b/package.json
index bf7281a..4b73959 100644
--- a/package.json
+++ b/package.json
@@ -10,7 +10,7 @@
   "author": "Intrinsic Technologies",
   "license": "MIT",
   "dependencies": {
-    "@axiom-crypto/client": "2.0.3",
+    "@axiom-crypto/client": "2.0.4",
     "dotenv": "^16.3.1",
     "ethers": "^6.10.0"
   },
diff --git a/test/AverageBalance.t.sol b/test/AverageBalance.t.sol
index 6c2526d..a20131b 100644
--- a/test/AverageBalance.t.sol
+++ b/test/AverageBalance.t.sol
@@ -1,41 +1,49 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.19;
 
-import { AxiomTest, AxiomVm } from "@axiom-crypto/v2-periphery/test/AxiomTest.sol";
-import { IAxiomV2Query } from "@axiom-crypto/v2-periphery/interfaces/query/IAxiomV2Query.sol";
+import "@axiom-crypto/axiom-std/AxiomTest.sol";
 
 import { AverageBalance } from "../src/AverageBalance.sol";
 
 contract AverageBalanceTest is AxiomTest {
+    using Axiom for Query;
+
+    struct AxiomInput {
+        uint64 blockNumber;
+        address addr;
+    }
+
     AverageBalance public averageBalance;
+    AxiomInput public input;
+    bytes32 public querySchema;
 
     function setUp() public {
         _createSelectForkAndSetupAxiom("sepolia", 5_103_100);
 
-        inputPath = "app/axiom/data/inputs.json";
-        querySchema = axiomVm.compile("app/axiom/average.circuit.ts", inputPath);
+        input =
+            AxiomInput({ blockNumber: 4_205_938, addr: address(0x8018fe32fCFd3d166E8b4c4E37105318A84BA11b) });
+        querySchema = axiomVm.readCircuit("app/axiom/average.circuit.ts");
         averageBalance = new AverageBalance(axiomV2QueryAddress, uint64(block.chainid), querySchema);
     }
 
-    function test_axiomSendQuery() public {
-        AxiomVm.AxiomSendQueryArgs memory args =
-            axiomVm.sendQueryArgs(inputPath, address(averageBalance), callbackExtraData, feeData);
-
-        axiomV2Query.sendQuery{ value: args.value }(
-            args.sourceChainId,
-            args.dataQueryHash,
-            args.computeQuery,
-            args.callback,
-            args.feeData,
-            args.userSalt,
-            args.refundee,
-            args.dataQuery
-        );
-    }
+    /// @dev Simple demonstration of testing an Axiom client contract using Axiom cheatcodes
+    function test_simple_example() public {
+        // create a query into Axiom with default parameters
+        Query memory q = query(querySchema, abi.encode(input), address(averageBalance));
+
+        // send the query to Axiom
+        q.send();
+
+        // prank fulfillment of the query, returning the Axiom results 
+        bytes32[] memory results = q.prankFulfill();
+
+        // parse Axiom results and verify length is as expected
+        assertEq(results.length, 3);
+        uint256 blockNumber = uint256(results[0]);
+        address addr = address(uint160(uint256(results[1])));
+        uint256 avg = uint256(results[2]);
 
-    function test_axiomCallback() public {
-        AxiomVm.AxiomFulfillCallbackArgs memory args =
-            axiomVm.fulfillCallbackArgs(inputPath, address(averageBalance), callbackExtraData, feeData, msg.sender);
-        axiomVm.prankCallback(args);
+        // verify the average balance recorded in AverageBalance is as expected
+        assertEq(avg, averageBalance.provenAverageBalances(blockNumber, addr));
     }
 }