add test cases and fix mypy issue

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",

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()