From fbffca50ea673320002e45f740ad737475702c93 Mon Sep 17 00:00:00 2001
From: Alireza Haghshenas <alirezahaghshenas@gmail.com>
Date: Thu, 28 Sep 2023 10:59:03 -0400
Subject: [PATCH] chore: apply suggestions from PR comments

---
 docs/staking.md                               |  2 +-
 integration-tests/contract-staking.spec.ts    | 65 ++++++++++++++++++-
 packages/taquito/src/contract/errors.ts       |  4 +-
 packages/taquito/src/contract/interface.ts    |  2 +-
 .../src/contract/rpc-contract-provider.ts     |  2 +-
 packages/taquito/src/operations/types.ts      |  9 +--
 .../contract/rpc-contract-provider.spec.ts    |  2 +-
 7 files changed, 71 insertions(+), 15 deletions(-)

diff --git a/docs/staking.md b/docs/staking.md
index 96ac2dab44..eaea83ec39 100644
--- a/docs/staking.md
+++ b/docs/staking.md
@@ -7,7 +7,7 @@ author: Alireza Haghshenas
 Since the first Tezos protocol, Tez owners could participate in baking, endorsement (now renamed to attestation), and voting by `delegating` their tez to bakers.
 Starting with Oxford, there is now an additional way of participating which is called `staking`. You can read more about it (here)[https://tezos.gitlab.io/protocols/018_oxford.html#adaptive-issuance-ongoing] and (here)[https://spotlight.tezos.com/announcing-oxford-tezos-15th-protocol-upgrade-proposal/] (under "Adaptive Issuance and Staking").
 
-In Taquito, we support three new operations: `stake`, `unstake`, and `finalize_unstake`.
+In Taquito, we support three new operations: `stake`, `unstake`, and `finalize_unstake` implemented as pseudo operations (entrypoints on implicit accounts).
 
 ## stake
 By `stake`ing your tez to a baker that accepts staking, your tez will be frozen. Also, unlike delegation, your tez is subject to slashing in case the baker misbehaves. In return, you can receive more rewards for baking and attestation (formerly endorsement).
diff --git a/integration-tests/contract-staking.spec.ts b/integration-tests/contract-staking.spec.ts
index 063733f182..c1d13939b1 100644
--- a/integration-tests/contract-staking.spec.ts
+++ b/integration-tests/contract-staking.spec.ts
@@ -4,14 +4,14 @@ import { CONFIGS, isSandbox, sleep } from "./config";
 CONFIGS().forEach(({ lib, rpc, setup, protocol }) => {
     const Tezos = lib;
     describe(`Test staking through contract API using: ${rpc}`, () => {
-      const flextesaOxford = isSandbox({rpc}) && protocol === Protocols.Proxford ? test : test.skip;
+      const flextesaOxford = isSandbox({rpc}) && protocol === Protocols.ProxfordS ? test : test.skip;
       beforeEach(async (done) => {
           await setup(true)
           done()
         });
         flextesaOxford('Should be able to stake', async (done) => {
           const op = await Tezos.contract.stake({
-            amount: 100,
+            amount: 100000,
           });
           await op.confirmation()
           expect(op.hash).toBeDefined();
@@ -23,7 +23,7 @@ CONFIGS().forEach(({ lib, rpc, setup, protocol }) => {
         });
         flextesaOxford('Should be able to unstake', async (done) => {
           const op = await Tezos.contract.unstake({
-            amount: 100,
+            amount: 100000,
           });
           await op.confirmation()
           expect(op.hash).toBeDefined();
@@ -39,5 +39,64 @@ CONFIGS().forEach(({ lib, rpc, setup, protocol }) => {
           expect(op.hash).toBeDefined();
           done();
         });
+        flextesaOxford.skip('Should fail with unstake amount too large, could not produce: gets proto.018-Proxford.tez.subtraction_underflow when the amount is too high', async (done) => {
+          expect(async () => {
+            const op = await Tezos.contract.stake({
+              amount: 1000000000000,
+            });
+            await op.confirmation()
+          }).rejects.toThrowError(/Invalid unstake request amount/);
+          done();
+        });
+        flextesaOxford('Should fail with invalid_nonzero_transaction_amount', async (done) => {
+          expect(async () => {
+            const address = await Tezos.signer.publicKeyHash();
+            const op = await Tezos.contract.transfer({
+              to: address,
+              amount: 10000,
+              parameter: {
+                entrypoint: "unstake",
+                value: {
+                  int: "10000",
+                },
+              },
+            });
+            await op.confirmation()
+          }).rejects.toThrowError(/\(permanent\) proto.\d{3}-\w+\.operations\.invalid_nonzero_transaction_amount/);
+          done();
+        });
+        flextesaOxford('Should fail with invalid_unstake_request_amount', async (done) => {
+          expect(async () => {
+            const address = await Tezos.signer.publicKeyHash();
+            const op = await Tezos.contract.transfer({
+              to: address,
+              amount: 0,
+              parameter: {
+                entrypoint: "unstake",
+                value: {
+                  int: "-10000",
+                },
+              },
+            });
+            await op.confirmation()
+          }).rejects.toThrowError(/\(permanent\) proto.\d{3}-\w+\.operations\.invalid_unstake_request_amount/);
+          done();
+        });
+        flextesaOxford('Should fail with invalid_self_transaction_destination', async (done) => {
+          expect(async () => {
+            const op = await Tezos.contract.transfer({
+              to: 'tz1Yju7jmmsaUiG9qQLoYv35v5pHgnWoLWbt',
+              amount: 0,
+              parameter: {
+                entrypoint: "unstake",
+                value: {
+                  int: "10000",
+                },
+              },
+            });
+            await op.confirmation()
+          }).rejects.toThrowError(/\(permanent\) proto.\d{3}-\w+\.operations\.invalid_self_transaction_destination/);
+          done();
+        });
     });
 });
diff --git a/packages/taquito/src/contract/errors.ts b/packages/taquito/src/contract/errors.ts
index ebf7fafd34..c9a6fbd22d 100644
--- a/packages/taquito/src/contract/errors.ts
+++ b/packages/taquito/src/contract/errors.ts
@@ -22,13 +22,13 @@ export class InvalidParameterError extends ParameterValidationError {
 
 /**
  *  @category Error
- *  @description Error that indicates an invalid delegation source contract address being passed or used
+ *  @description Error that indicates an invalid delegation source originated address being passed or used
  */
 export class InvalidDelegationSource extends ParameterValidationError {
   constructor(public readonly source: string) {
     super();
     this.name = `InvalidDelegationSource`;
-    this.message = `Since Babylon delegation source can no longer be a contract address ${source}. Please use the smart contract abstraction to set your delegate.`;
+    this.message = `Since Babylon delegation source can no longer be an originated address ${source}. Please use the smart contract abstraction to set your delegate.`;
   }
 }
 
diff --git a/packages/taquito/src/contract/interface.ts b/packages/taquito/src/contract/interface.ts
index 8db8256ca9..ee4b570cde 100644
--- a/packages/taquito/src/contract/interface.ts
+++ b/packages/taquito/src/contract/interface.ts
@@ -299,7 +299,7 @@ export interface ContractProvider extends StorageProvider {
 
   /**
    *
-   * @description Unstake the given amount. If "everything" is given as amount, unstakes everything from the staking balance. Unstaked tez remains frozen for a set amount of cycles (the slashing period) after the operation. Once this period is over, the operation "finalize unstake" must be called for the funds to appear in the liquid balance. Will sign and inject an operation using the current context
+   * @description Unstake the given amount. If `amount` is larger than the currently staked amount, unstakes everything from the staking balance. Unstaked tez remains frozen for a set amount of cycles (the slashing period) after the operation. Once this period is over, the operation "finalize unstake" must be called for the funds to appear in the liquid balance. Will sign and inject an operation using the current context
    *
    * @returns An operation handle with the result from the rpc node
    *
diff --git a/packages/taquito/src/contract/rpc-contract-provider.ts b/packages/taquito/src/contract/rpc-contract-provider.ts
index c3e6a2a39a..79ba12a05f 100644
--- a/packages/taquito/src/contract/rpc-contract-provider.ts
+++ b/packages/taquito/src/contract/rpc-contract-provider.ts
@@ -792,7 +792,7 @@ export class RpcContractProvider extends Provider implements ContractProvider, S
 
   /**
    *
-   * @description Unstake the given amount. If "everything" is given as amount, unstakes everything from the staking balance. Unstaked tez remains frozen for a set amount of cycles (the slashing period) after the operation. Once this period is over, the operation "finalize unstake" must be called for the funds to appear in the liquid balance.
+   * @description Unstake the given amount. If `amount` is larger than the currently staked amount, unstakes everything from the staking balance. Unstaked tez remains frozen for a set amount of cycles (the slashing period) after the operation. Once this period is over, the operation "finalize unstake" must be called for the funds to appear in the liquid balance.
    *
    * @returns An operation handle with the result from the rpc node
    *
diff --git a/packages/taquito/src/operations/types.ts b/packages/taquito/src/operations/types.ts
index 64b92cc858..ff4b3bf438 100644
--- a/packages/taquito/src/operations/types.ts
+++ b/packages/taquito/src/operations/types.ts
@@ -9,8 +9,7 @@ import {
 } from '@taquito/rpc';
 import { BlockIdentifier } from '../read-provider/interface';
 import { InvalidAddressError, InvalidAmountError } from '@taquito/core';
-import { ValidationResult, invalidDetail, validateAddress } from '@taquito/utils';
-import { InvalidStakingSource } from '../contract/errors';
+import { ValidationResult, invalidDetail, validateKeyHash } from '@taquito/utils';
 
 export { OpKind } from '@taquito/rpc';
 
@@ -576,6 +575,7 @@ export interface InternalStakingParams {
   parameter: TransactionOperationParameter;
 }
 
+// Helper type to make a field optional
 type PartialBy<T, K extends keyof T> = Omit<T, K> & Partial<Pick<T, K>>;
 export type StakingParams = PartialBy<InternalStakingParams, 'source'>;
 export type StakeParams = Omit<StakingParams, 'parameter'>;
@@ -601,13 +601,10 @@ export const validateStakingParams = (
     }
   }
   if (params.source !== undefined) {
-    const sourceValidation = validateAddress(params.source);
+    const sourceValidation = validateKeyHash(params.source);
     if (sourceValidation !== ValidationResult.VALID) {
       throw new InvalidAddressError(params.source, invalidDetail(sourceValidation));
     }
-    if (/kt/i.test(params.source)) {
-      throw new InvalidStakingSource(params.source);
-    }
   }
 };
 
diff --git a/packages/taquito/test/contract/rpc-contract-provider.spec.ts b/packages/taquito/test/contract/rpc-contract-provider.spec.ts
index 96ecaf2d24..1caa63acc3 100644
--- a/packages/taquito/test/contract/rpc-contract-provider.spec.ts
+++ b/packages/taquito/test/contract/rpc-contract-provider.spec.ts
@@ -26,7 +26,7 @@ import { OpKind, ParamsWithKind, TransferTicketParams } from '../../src/operatio
 import { NoopParser } from '../../src/taquito';
 import { OperationBatch } from '../../src/batch/rpc-batch-provider';
 import { PvmKind } from '@taquito/rpc';
-import { InvalidAddressError } from '@taquito/utils';
+import { InvalidAddressError } from '@taquito/core';
 import { InvalidAmountError } from '@taquito/core';
 
 /**