Starknet: enum in map

lib.cairo

#[derive(Drop, Serde, Copy, PartialEq, Debug, starknet::Store)]
pub enum Choice {
    STRK,
    BTC,
}

#[starknet::interface]
pub trait IHelloStarknet<TContractState> {
    fn create_elem(ref self: TContractState, choice: Choice) -> u256;
    fn get_elem(self: @TContractState, elem_id: u256) -> Choice;
}

#[starknet::contract]
mod HelloStarknet {
    use core::starknet::storage::{
        StoragePointerReadAccess, StoragePointerWriteAccess, Map, StoragePathEntry,
    };
    use super::Choice;

    #[storage]
    struct Storage {
        choices: Map<u256, Choice>,
        choices_count: u256,
    }

    #[constructor]
    fn constructor(ref self: ContractState) {
        self.choices_count.write(0);
    }

    #[abi(embed_v0)]
    impl HelloStarknetImpl of super::IHelloStarknet<ContractState> {
        fn create_elem(ref self: ContractState, choice: Choice) -> u256 {
            let elem_id = self.choices_count.read() + 1;
            self.choices_count.write(elem_id);
            self.choices.entry(elem_id).write(choice);
            elem_id
        }

        fn get_elem(self: @ContractState, elem_id: u256) -> Choice {
            self.choices.entry(elem_id).read()
        }
    }
}

test_contract.cairo

use starknet::ContractAddress;

use snforge_std::{declare, ContractClassTrait, DeclareResultTrait};

use enum_map::IHelloStarknetDispatcher;
use enum_map::IHelloStarknetDispatcherTrait;
use enum_map::Choice;

fn deploy_contract(name: ByteArray) -> ContractAddress {
    let contract = declare(name).unwrap().contract_class();
    let (contract_address, _) = contract.deploy(@ArrayTrait::new()).unwrap();
    contract_address
}

#[test]
fn test_get_valid_elem() {
    let contract_address = deploy_contract("HelloStarknet");

    let dispatcher = IHelloStarknetDispatcher { contract_address };
    let element_id = dispatcher.create_elem(Choice::STRK);

    let choice = dispatcher.get_elem(element_id);
    assert_eq!(choice, Choice::STRK, "get_elem failed");
}

#[test]
fn test_get_invalid_elem() {
    let contract_address = deploy_contract("HelloStarknet");

    let dispatcher = IHelloStarknetDispatcher { contract_address };
    let element_id = dispatcher.create_elem(Choice::STRK) + 1;

    let choice = dispatcher.get_elem(element_id);
    assert_eq!(choice, Choice::STRK, "Should failed before");
}