From c3e648834496d5ef4ca0960d3ce8306831c4d6d7 Mon Sep 17 00:00:00 2001
From: ejMina226 <118474890+ejMina226@users.noreply.github.com>
Date: Tue, 7 Jan 2025 18:20:59 +0000
Subject: [PATCH] Fix tests

---
 packages/module/test/modules/Balances.test.ts |  8 ++++----
 packages/protocol/test/BlockProver.test.ts    | 10 +++++-----
 packages/protocol/test/TestingProtocol.ts     | 15 +++++++++++----
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/packages/module/test/modules/Balances.test.ts b/packages/module/test/modules/Balances.test.ts
index a425c18fe..b0b9fc4fa 100644
--- a/packages/module/test/modules/Balances.test.ts
+++ b/packages/module/test/modules/Balances.test.ts
@@ -68,7 +68,7 @@ describe("balances", () => {
     it("should compile and prove a method execution", async () => {
       expect.assertions(3);
 
-      runtime.zkProgrammable.appChain?.setProofsEnabled(true);
+      runtime.appChain?.setProofsEnabled(true);
 
       const executionContext = container.resolve(RuntimeMethodExecutionContext);
       executionContext.setup({
@@ -80,7 +80,7 @@ describe("balances", () => {
         "1439144406936083177718146178121957896974210157062549589517697792374542035761";
       const expectedStatus = true;
 
-      await runtime.zkProgrammable.zkProgram[0].compile();
+      await runtime.zkProgram[0].compile();
 
       await balances.getTotalSupply();
 
@@ -88,9 +88,9 @@ describe("balances", () => {
 
       const proof = await result.prove<Proof<undefined, MethodPublicOutput>>();
 
-      const verified = await runtime.zkProgrammable.zkProgram[0].verify(proof);
+      const verified = await runtime.zkProgram[0].verify(proof);
 
-      runtime.zkProgrammable.appChain?.setProofsEnabled(false);
+      runtime.appChain?.setProofsEnabled(false);
 
       expect(verified).toBe(true);
 
diff --git a/packages/protocol/test/BlockProver.test.ts b/packages/protocol/test/BlockProver.test.ts
index b02dd6909..644c96c2a 100644
--- a/packages/protocol/test/BlockProver.test.ts
+++ b/packages/protocol/test/BlockProver.test.ts
@@ -2,7 +2,7 @@
 import {
   AreProofsEnabled,
   PlainZkProgram,
-  ZkProgrammable,
+  ZkProgramFactory,
 } from "@proto-kit/common";
 import { Bool, Field, Proof, UInt64, ZkProgram } from "o1js";
 import "reflect-metadata";
@@ -30,10 +30,9 @@ class MockAppChain implements AreProofsEnabled {
   }
 }
 
-class RuntimeZkProgrammable extends ZkProgrammable<
-  undefined,
-  MethodPublicOutput
-> {
+class RuntimeZkProgrammable
+  implements ZkProgramFactory<undefined, MethodPublicOutput>
+{
   get appChain(): AreProofsEnabled | undefined {
     return new MockAppChain();
   }
@@ -51,6 +50,7 @@ class RuntimeZkProgrammable extends ZkProgrammable<
         verify: program.verify.bind(program),
         analyzeMethods: program.analyzeMethods.bind(program),
         methods: {},
+        proofsEnabled: true,
         Proof: ZkProgram.Proof(program),
       },
     ];
diff --git a/packages/protocol/test/TestingProtocol.ts b/packages/protocol/test/TestingProtocol.ts
index 09e82cbd1..548d23cf2 100644
--- a/packages/protocol/test/TestingProtocol.ts
+++ b/packages/protocol/test/TestingProtocol.ts
@@ -1,4 +1,8 @@
-import { WithZkProgrammable, ZkProgrammable } from "@proto-kit/common";
+import {
+  PlainZkProgram,
+  WithZkProgram,
+  ZkProgramFactory,
+} from "@proto-kit/common";
 import { container } from "tsyringe";
 
 import {
@@ -11,9 +15,12 @@ import {
   StateTransitionProver,
 } from "../src";
 
-class RuntimeMock implements WithZkProgrammable<undefined, MethodPublicOutput> {
-  zkProgrammable: ZkProgrammable<undefined, MethodPublicOutput> =
-    undefined as unknown as ZkProgrammable<undefined, MethodPublicOutput>;
+class RuntimeMock implements WithZkProgram<undefined, MethodPublicOutput> {
+  zkProgramFactory: ZkProgramFactory<undefined, MethodPublicOutput> =
+    undefined as unknown as ZkProgramFactory<undefined, MethodPublicOutput>;
+
+  zkProgram: PlainZkProgram<undefined, MethodPublicOutput>[] =
+    undefined as unknown as PlainZkProgram<undefined, MethodPublicOutput>[];
 }
 
 export function createAndInitTestingProtocol() {