From 77eff458125828565ceb35d953d63c52ac5f0566 Mon Sep 17 00:00:00 2001
From: FabijanC <fc@spaceshard.io>
Date: Mon, 4 Nov 2024 12:42:44 +0100
Subject: [PATCH] Implement unsubscription TODOs [skip ci]

---
 .../src/api/json_rpc/endpoints_ws.rs          |  2 +-
 .../src/api/json_rpc/error.rs                 |  7 ++++++
 .../src/api/json_rpc/mod.rs                   |  4 ++--
 .../starknet-devnet-server/src/subscribe.rs   | 23 +++++++++++++++----
 4 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/crates/starknet-devnet-server/src/api/json_rpc/endpoints_ws.rs b/crates/starknet-devnet-server/src/api/json_rpc/endpoints_ws.rs
index c823eec91..a3a61c0e8 100644
--- a/crates/starknet-devnet-server/src/api/json_rpc/endpoints_ws.rs
+++ b/crates/starknet-devnet-server/src/api/json_rpc/endpoints_ws.rs
@@ -30,7 +30,7 @@ impl JsonRpcHandler {
                     }),
                 )?;
 
-                socket_context.unsubscribe(subscription_id).await;
+                socket_context.unsubscribe(rpc_request_id, subscription_id).await?;
                 Ok(())
             }
         }
diff --git a/crates/starknet-devnet-server/src/api/json_rpc/error.rs b/crates/starknet-devnet-server/src/api/json_rpc/error.rs
index 2c1eb6b8c..b5ecaea69 100644
--- a/crates/starknet-devnet-server/src/api/json_rpc/error.rs
+++ b/crates/starknet-devnet-server/src/api/json_rpc/error.rs
@@ -64,6 +64,8 @@ pub enum ApiError {
     CompiledClassHashMismatch,
     #[error("Cannot go back more than 1024 blocks")]
     TooManyBlocksBack,
+    #[error("Invalid subscription id")]
+    InvalidSubscriptionId,
 }
 
 impl ApiError {
@@ -212,6 +214,11 @@ impl ApiError {
                 message: error_message.into(),
                 data: None,
             },
+            ApiError::InvalidSubscriptionId => RpcError {
+                code: crate::rpc_core::error::ErrorCode::ServerError(66),
+                message: error_message.into(),
+                data: None,
+            },
         }
     }
 }
diff --git a/crates/starknet-devnet-server/src/api/json_rpc/mod.rs b/crates/starknet-devnet-server/src/api/json_rpc/mod.rs
index 8ae35e375..e71fdb17b 100644
--- a/crates/starknet-devnet-server/src/api/json_rpc/mod.rs
+++ b/crates/starknet-devnet-server/src/api/json_rpc/mod.rs
@@ -1423,8 +1423,8 @@ mod response_tests {
     use crate::api::json_rpc::ToRpcResponseResult;
 
     #[test]
-    fn serializing_starknet_response_empty_variant_has_to_produce_empty_json_object_when_converted_to_rpc_result(
-    ) {
+    fn serializing_starknet_response_empty_variant_has_to_produce_empty_json_object_when_converted_to_rpc_result()
+     {
         assert_eq!(
             r#"{"result":{}}"#,
             serde_json::to_string(
diff --git a/crates/starknet-devnet-server/src/subscribe.rs b/crates/starknet-devnet-server/src/subscribe.rs
index 67740f27e..917eb2f11 100644
--- a/crates/starknet-devnet-server/src/subscribe.rs
+++ b/crates/starknet-devnet-server/src/subscribe.rs
@@ -1,4 +1,5 @@
-use std::{collections::HashMap, sync::Arc};
+use std::collections::HashMap;
+use std::sync::Arc;
 
 use axum::extract::ws::{Message, WebSocket};
 use futures::stream::SplitSink;
@@ -7,6 +8,7 @@ use serde::{self, Serialize};
 use starknet_types::rpc::block::BlockHeader;
 use tokio::sync::Mutex;
 
+use crate::api::json_rpc::error::ApiError;
 use crate::rpc_core::request::Id;
 
 pub type SocketId = u64;
@@ -118,10 +120,21 @@ impl SocketContext {
         subscription_id
     }
 
-    pub async fn unsubscribe(&mut self, subscription_id: SubscriptionId) {
+    pub async fn unsubscribe(
+        &mut self,
+        rpc_request_id: Id,
+        subscription_id: SubscriptionId,
+    ) -> Result<(), ApiError> {
         match self.subscriptions.remove(&subscription_id) {
-            Some(_) => todo!("return true"),
-            None => todo!("return INVALID_SUBSCRIPTION_ID"),
+            Some(_) => {
+                self.send(SubscriptionResponse::Confirmation {
+                    rpc_request_id,
+                    result: SubscriptionConfirmation::UnsubscriptionConfirmation(true),
+                })
+                .await;
+                Ok(())
+            }
+            None => Err(ApiError::InvalidSubscriptionId),
         }
     }
 
@@ -137,7 +150,7 @@ impl SocketContext {
                         self.notify(*subscription_id, data.clone()).await;
                     }
                 }
-                other => println!("DEBUG unsupported subscription: {other:?}"),
+                other => panic!("DEBUG unsupported subscription: {other:?}"),
             }
         }
     }