feat: add parties reset and re-inclusion mechanism

Author: arobsn

.changeset/hip-falcons-take.md

@@ -0,0 +1,5 @@
+---
+"@fleet-sdk/mock-chain": patch
+---
+
+Add parties cleanup and re-inclusion

packages/mock-chain/src/mockChain.spec.ts

@@ -95,12 +95,47 @@ describe("Mock chain instantiation", () => {
     chain.reset();
 
     // should clear and reset state to original values
+    expect(chain.parties).to.have.length(2); // parties should not be cleared by default
     expect(chain.height).to.be.equal(creationHeight);
     expect(chain.timestamp).to.be.equal(creationTimestamp);
     expect(bob.utxos).to.have.length(0);
     expect(alice.utxos).to.have.length(0);
   });
 
+  it("Should reset the chain and clean up the parties", () => {
+    const chain = new MockChain(234897);
+    const creationHeight = chain.height;
+    const creationTimestamp = chain.timestamp;
+
+    const bob = chain
+      .newParty()
+      .withBalance({ nanoergs: 1000000n })
+      .withBalance({ nanoergs: 19827398237n });
+    const alice = chain.newParty().withBalance({ nanoergs: 6000000n });
+
+    expect(bob.utxos).to.have.length(2);
+    expect(alice.utxos).to.have.length(1);
+
+    chain.newBlocks(100);
+
+    expect(chain.height).to.be.greaterThan(creationHeight);
+    expect(chain.timestamp).to.be.greaterThan(creationTimestamp);
+
+    // reset the chain
+    chain.reset({ clearParties: true });
+
+    // should clear and reset state to original values
+    expect(chain.parties).to.have.length(0); // parties should be cleared
+    expect(chain.height).to.be.equal(creationHeight);
+    expect(chain.timestamp).to.be.equal(creationTimestamp);
+    expect(bob.utxos).to.have.length(0);
+    expect(alice.utxos).to.have.length(0);
+
+    bob.addBalance({ nanoergs: 1000000n });
+    expect(bob.utxos).to.have.length(1); // parties should be re-added after adding balance
+    expect(chain.parties).to.have.length(1);
+  });
+
   it("Should simulate new blocks", () => {
     const blockTime = 120_000;
     const height = 100;

packages/mock-chain/src/mockChain.ts

@@ -42,6 +42,10 @@ export type BlockState = {
   parameters: BlockchainParameters;
 };
 
+export type ResetOptions = {
+  clearParties?: boolean;
+};
+
 export class MockChain {
   readonly #parties: MockChainParty[];
   readonly #tip: BlockState;
@@ -104,10 +108,14 @@ export class MockChain {
     }
   }
 
-  reset() {
+  reset(options?: ResetOptions) {
     this.clearUTxOSet();
     this.#tip.height = this.#base.height;
     this.#tip.timestamp = this.#base.timestamp;
+
+    if (options?.clearParties) {
+      this.#parties.length = 0;
+    }
   }
 
   newParty(name?: string): KeyedMockChainParty;
@@ -120,8 +128,12 @@ export class MockChain {
     );
   }
 
-  addParty(ergoTree: HexString, name?: string): NonKeyedMockChainParty {
-    return this.#pushParty(new NonKeyedMockChainParty(this, ergoTree, name));
+  addParty(contractOrParty: HexString | MockChainParty, name?: string): NonKeyedMockChainParty {
+    return this.#pushParty(
+      typeof contractOrParty === "string"
+        ? new NonKeyedMockChainParty(this, contractOrParty, name)
+        : contractOrParty
+    );
   }
 
   #pushParty<T extends MockChainParty>(party: T): T {

packages/mock-chain/src/party/mockChainParty.ts

@@ -68,6 +68,11 @@ export abstract class MockChainParty {
   addUTxOs(utxos: MockUTxOInput): MockChainParty {
     this._utxos.add(utxos);
 
+    // Ensure the party is registered in the chain
+    if (!this._chain.parties.includes(this)) {
+      this._chain.addParty(this);
+    }
+
     return this;
   }