Add tests for withdraw functionality in StreamContract

contracts/stream_contract/src/lib.rs

@@ -1,7 +1,18 @@
 #![no_std]
-use soroban_sdk::{contract, contractimpl, contracttype, contracterror, Address, Env, Symbol, symbol_short, token};
 
-#[derive(Clone)]
+use soroban_sdk::{
+    contract, contracterror, contractimpl, contracttype, panic_with_error, token, Address, Env,
+    Symbol,
+};
+
+#[contracttype]
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub enum DataKey {
+    Stream(u64),
+    StreamCounter,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
 #[contracttype]
 pub struct Stream {
     pub sender: Address,
@@ -24,7 +35,6 @@ pub enum StreamError {
     StreamInactive = 4,
 }
 
-// Event definitions for indexing
 #[contracttype]
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct StreamCreatedEvent {
@@ -77,26 +87,48 @@ impl StreamContract {
         duration: u64,
     ) -> u64 {
         sender.require_auth();
-        // Placeholder for stream creation logic
-        // 1. Transfer tokens to contract
-        // 2. Store stream state
 
-        // Generate stream ID (placeholder - use proper counter in production)
-        let stream_id: u64 = env.ledger().sequence() as u64;
+        if amount <= 0 || duration == 0 {
+            panic_with_error!(&env, StreamError::InvalidAmount);
+        }
+
+        let stream_id = Self::get_next_stream_id(&env);
         let start_time = env.ledger().timestamp();
+        let rate_per_second = amount / (duration as i128);
+
+        let token_client = token::Client::new(&env, &token_address);
+        let contract_address = env.current_contract_address();
+        token_client.transfer(&sender, &contract_address, &amount);
+
+        let stream = Stream {
+            sender: sender.clone(),
+            recipient: recipient.clone(),
+            token_address: token_address.clone(),
+            rate_per_second,
+            deposited_amount: amount,
+            withdrawn_amount: 0,
+            start_time,
+            last_update_time: start_time,
+            is_active: true,
+        };
+
+        env.storage()
+            .persistent()
+            .set(&DataKey::Stream(stream_id), &stream);
 
-        // Emit StreamCreated event
         env.events().publish(
             (Symbol::new(&env, "stream_created"), stream_id),
             StreamCreatedEvent {
                 stream_id,
-                sender: sender.clone(),
-                recipient: recipient.clone(),
-                rate,
-                token_address: token_address.clone(),
+                sender,
+                recipient,
+                rate: rate_per_second,
+                token_address,
                 start_time,
-            }
+            },
         );
+
+        stream_id
     }
 
     fn get_next_stream_id(env: &Env) -> u64 {
@@ -112,158 +144,144 @@ impl StreamContract {
         next_id
     }
 
-    pub fn withdraw(_env: Env, _recipient: Address, _stream_id: u64) {
-        // Placeholder for withdraw logic
-        // 1. Calculate claimable amount based on time delta
-        // 2. Transfer tokens to recipient
-        // 3. Update stream state
+    pub fn withdraw(env: Env, recipient: Address, stream_id: u64) -> Result<i128, StreamError> {
+        recipient.require_auth();
 
-        // Placeholder amount calculation
-        let amount: i128 = 0; // Calculate actual amount in production
-        let timestamp = env.ledger().timestamp();
+        let storage = env.storage().persistent();
+        let stream_key = DataKey::Stream(stream_id);
 
-        // Emit TokensWithdrawn event
-        env.events().publish(
-            (Symbol::new(&env, "tokens_withdrawn"), stream_id),
-            TokensWithdrawnEvent {
-                stream_id,
-                recipient: recipient.clone(),
-                amount,
-                timestamp,
-            }
-        );
-    }
+        let mut stream: Stream = match storage.get(&stream_key) {
+            Some(s) => s,
+            None => return Err(StreamError::StreamNotFound),
+        };
 
-    pub fn cancel_stream(env: Env, sender: Address, stream_id: u64) {
-        sender.require_auth();
-        // Placeholder for cancel logic
-        // 1. Calculate amount already withdrawn
-        // 2. Return remaining tokens to sender
-        // 3. Mark stream as cancelled
+        if stream.recipient != recipient {
+            return Err(StreamError::Unauthorized);
+        }
 
-        // Placeholder values
-        let recipient = sender.clone(); // Get actual recipient from storage in production
-        let amount_withdrawn: i128 = 0; // Calculate actual amount in production
+        if !stream.is_active {
+            return Err(StreamError::StreamInactive);
+        }
+
+        let claimable = stream.deposited_amount - stream.withdrawn_amount;
+        if claimable <= 0 {
+            return Err(StreamError::InvalidAmount);
+        }
+
+        let token_client = token::Client::new(&env, &stream.token_address);
+        let contract_address = env.current_contract_address();
+        token_client.transfer(&contract_address, &recipient, &claimable);
+
+        stream.withdrawn_amount += claimable;
+        stream.last_update_time = env.ledger().timestamp();
+        if stream.withdrawn_amount >= stream.deposited_amount {
+            stream.is_active = false;
+        }
+        storage.set(&stream_key, &stream);
 
-        // Emit StreamCancelled event
         env.events().publish(
-            (Symbol::new(&env, "stream_cancelled"), stream_id),
-            StreamCancelledEvent {
+            (Symbol::new(&env, "tokens_withdrawn"), stream_id),
+            TokensWithdrawnEvent {
                 stream_id,
-                sender: sender.clone(),
                 recipient,
-                amount_withdrawn,
-            }
+                amount: claimable,
+                timestamp: stream.last_update_time,
+            },
         );
+
+        Ok(claimable)
     }
 
-    /// Allows the sender to add more funds to an existing stream
-    /// This extends the duration of the stream without creating a new one
-    pub fn top_up_stream(env: Env, sender: Address, stream_id: u64, amount: i128) -> Result<(), StreamError> {
-        // Require sender authentication
+    pub fn cancel_stream(env: Env, sender: Address, stream_id: u64) -> Result<(), StreamError> {
         sender.require_auth();
 
-        // Validate amount is positive
-        if amount <= 0 {
-            return Err(StreamError::InvalidAmount);
-        }
-
-        // Get the stream from storage
         let storage = env.storage().persistent();
-        let stream_key = (symbol_short!("STREAMS"), stream_id);
+        let stream_key = DataKey::Stream(stream_id);
 
         let mut stream: Stream = match storage.get(&stream_key) {
             Some(s) => s,
             None => return Err(StreamError::StreamNotFound),
         };
 
-        // Verify the caller is the original sender
         if stream.sender != sender {
             return Err(StreamError::Unauthorized);
         }
 
-        // Verify stream is still active
         if !stream.is_active {
             return Err(StreamError::StreamInactive);
         }
 
-        // Transfer tokens from sender to contract
-        let token_client = token::Client::new(&env, &stream.token_address);
-        let contract_address = env.current_contract_address();
-        token_client.transfer(&sender, &contract_address, &amount);
-
-        // Update stream state with additional deposit
-        stream.deposited_amount += amount;
+        stream.is_active = false;
         stream.last_update_time = env.ledger().timestamp();
-
-        // Save updated stream back to storage
+        let recipient = stream.recipient.clone();
+        let amount_withdrawn = stream.withdrawn_amount;
         storage.set(&stream_key, &stream);
 
-        // Emit StreamToppedUp event
         env.events().publish(
-            (Symbol::new(&env, "stream_topped_up"), stream_id),
-            StreamToppedUpEvent {
+            (Symbol::new(&env, "stream_cancelled"), stream_id),
+            StreamCancelledEvent {
                 stream_id,
-                sender: sender.clone(),
-                amount,
-                new_deposited_amount: stream.deposited_amount,
-            }
+                sender,
+                recipient,
+                amount_withdrawn,
+            },
         );
 
         Ok(())
     }
 
-    /// Allows the sender to add more funds to an existing stream
-    /// This extends the duration of the stream without creating a new one
-    pub fn top_up_stream(env: Env, sender: Address, stream_id: u64, amount: i128) -> Result<(), StreamError> {
-        // Require sender authentication
+    pub fn top_up_stream(
+        env: Env,
+        sender: Address,
+        stream_id: u64,
+        amount: i128,
+    ) -> Result<(), StreamError> {
         sender.require_auth();
 
-        // Validate amount is positive
         if amount <= 0 {
             return Err(StreamError::InvalidAmount);
         }
 
-        // Get the stream from storage
         let storage = env.storage().persistent();
-        let stream_key = (symbol_short!("STREAMS"), stream_id);
+        let stream_key = DataKey::Stream(stream_id);
 
         let mut stream: Stream = match storage.get(&stream_key) {
             Some(s) => s,
             None => return Err(StreamError::StreamNotFound),
         };
 
-        // Verify the caller is the original sender
         if stream.sender != sender {
             return Err(StreamError::Unauthorized);
         }
 
-        // Verify stream is still active
         if !stream.is_active {
             return Err(StreamError::StreamInactive);
         }
 
-        // Transfer tokens from sender to contract
         let token_client = token::Client::new(&env, &stream.token_address);
         let contract_address = env.current_contract_address();
         token_client.transfer(&sender, &contract_address, &amount);
 
-        // Update stream state with additional deposit
         stream.deposited_amount += amount;
         stream.last_update_time = env.ledger().timestamp();
 
-        // Save updated stream back to storage
         storage.set(&stream_key, &stream);
 
-        Ok(())
-    }
+        env.events().publish(
+            (Symbol::new(&env, "stream_topped_up"), stream_id),
+            StreamToppedUpEvent {
+                stream_id,
+                sender,
+                amount,
+                new_deposited_amount: stream.deposited_amount,
+            },
+        );
 
-    pub fn get_stream(env: Env, stream_id: u64) -> Option<Stream> {
-        env.storage().instance().get(&DataKey::Stream(stream_id))
+        Ok(())
     }
 
     pub fn get_stream(env: Env, stream_id: u64) -> Option<Stream> {
-        env.storage().instance().get(&DataKey::Stream(stream_id))
+        env.storage().persistent().get(&DataKey::Stream(stream_id))
     }
 }