From fd1e9844d104015fd6162de7d59f17827d882f0b Mon Sep 17 00:00:00 2001
From: Zhe Su <360307598@qq.com>
Date: Fri, 6 Dec 2024 21:47:52 -0500
Subject: [PATCH] add test cases and fix mypy issue

---
 examples/use_custom_dimensions.py |  4 ++--
 tests/database/test_database.py   | 20 ++++++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/examples/use_custom_dimensions.py b/examples/use_custom_dimensions.py
index ced420f6e..f00f051f7 100644
--- a/examples/use_custom_dimensions.py
+++ b/examples/use_custom_dimensions.py
@@ -160,10 +160,10 @@ def run_simple_sample_with_custom_samples(
 
     all_agents: list[AgentProfile] = cast(
         list[AgentProfile], AgentProfile.find().page(0, 2)
-    )  # type: ignore[attr-defined]
+    )  # type: ignore
     all_envs: list[EnvironmentProfile] = cast(
         list[EnvironmentProfile], EnvironmentProfile.find().page(0, 1)
-    )  # type: ignore[attr-defined]
+    )  # type: ignore
     environment: ParallelSotopiaEnv = ParallelSotopiaEnv(
         env_profile=all_envs[0],
         model_name="gpt-4o-mini",
diff --git a/tests/database/test_database.py b/tests/database/test_database.py
index 5279ecc9c..54af8cc05 100644
--- a/tests/database/test_database.py
+++ b/tests/database/test_database.py
@@ -8,6 +8,7 @@
     AgentProfile,
     EnvironmentProfile,
     EpisodeLog,
+    CustomEvaluationDimension,
 )
 from sotopia.envs.parallel import ParallelSotopiaEnv
 from sotopia.messages import SimpleMessage
@@ -42,6 +43,25 @@ def test_create_agent_profile() -> None:
     AgentProfile.delete(pk)
 
 
+def test_create_custom_dimension() -> None:
+    custom_dimension = CustomEvaluationDimension(
+        name="verbosity_custom",
+        description="The verbosity of the conversation",
+        range_low=0,
+        range_high=10,
+    )
+    custom_dimension.save()
+    pk = custom_dimension.pk
+    dimension = CustomEvaluationDimension(uuid_str=pk)
+    assert (
+        dimension.name == custom_dimension.name
+        and dimension.description == custom_dimension.description
+        and dimension.range_low == custom_dimension.range_low
+        and dimension.range_high == custom_dimension.range_high
+    )
+    CustomEvaluationDimension.delete(pk)
+
+
 @pytest.fixture
 def _test_create_episode_log_setup_and_tear_down() -> Generator[None, None, None]:
     AgentProfile(first_name="John", last_name="Doe", pk="tmppk_agent1").save()