Skip to content

Commit

Permalink
test(starknet_class_manager_types): add MemoryClassManagerClient for …
Browse files Browse the repository at this point in the history
…testing
  • Loading branch information
noamsp-starkware committed Jan 16, 2025
1 parent 3fc2de3 commit 3aba1f9
Show file tree
Hide file tree
Showing 3 changed files with 100 additions and 1 deletion.
3 changes: 3 additions & 0 deletions crates/starknet_class_manager_types/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@ license.workspace = true
repository.workspace = true
version.workspace = true

[features]
testing = []

[lints]
workspace = true

Expand Down
3 changes: 2 additions & 1 deletion crates/starknet_class_manager_types/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
#[cfg(any(feature = "testing", test))]
pub mod test_utils;
pub mod transaction_converter;

use std::sync::Arc;

use async_trait::async_trait;
Expand Down
95 changes: 95 additions & 0 deletions crates/starknet_class_manager_types/src/test_utils.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
use core::panic;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

use async_trait::async_trait;
use starknet_api::core::{ClassHash, CompiledClassHash};
use starknet_api::deprecated_contract_class::ContractClass as DeprecatedClass;
use starknet_api::state::SierraContractClass;

use crate::{
Class,
ClassId,
ClassManagerClient,
ClassManagerClientResult,
ClassManagerError,
ExecutableClass,
ExecutableClassHash,
};

pub struct MemoryClassManagerClient {
sierras: Arc<Mutex<HashMap<ClassHash, SierraContractClass>>>,
contracts: Arc<Mutex<HashMap<ClassHash, ExecutableClass>>>,
}

impl MemoryClassManagerClient {
pub fn new() -> Self {
Self {
sierras: Arc::new(Mutex::new(HashMap::new())),
contracts: Arc::new(Mutex::new(HashMap::new())),
}
}
}

impl Default for MemoryClassManagerClient {
fn default() -> Self {
Self::new()
}
}

#[async_trait]
impl ClassManagerClient for MemoryClassManagerClient {
async fn add_class(
&self,
class_id: ClassId,
class: Class,
) -> ClassManagerClientResult<ExecutableClassHash> {
if self.sierras.lock().unwrap().insert(class_id, class).is_some() {
panic!("Class already exists");
}

Ok(CompiledClassHash(class_id.0))
}

async fn get_executable(&self, class_id: ClassId) -> ClassManagerClientResult<ExecutableClass> {
let class = self
.contracts
.lock()
.unwrap()
.get(&class_id)
.cloned()
.ok_or_else(|| ClassManagerError::ClassNotFound { class_id })?;

Ok(class)
}

async fn get_sierra(&self, class_id: ClassId) -> ClassManagerClientResult<Class> {
let class = self
.sierras
.lock()
.unwrap()
.get(&class_id)
.cloned()
.ok_or_else(|| ClassManagerError::ClassNotFound { class_id })?;

Ok(class)
}

async fn add_deprecated_class(
&self,
class_id: ClassId,
class: DeprecatedClass,
) -> ClassManagerClientResult<()> {
if self
.contracts
.lock()
.unwrap()
.insert(class_id, starknet_api::contract_class::ContractClass::V0(class))
.is_some()
{
panic!("Class already exists");
}

Ok(())
}
}

0 comments on commit 3aba1f9

Please sign in to comment.