[fix] Update mock method names in tests

UKPLab · Jul 12, 2024 · 0efabc8 · 0efabc8
1 parent 69d6ddb
commit 0efabc8
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 19 deletions.
diff --git a/llm_roleplay/VERSION b/llm_roleplay/VERSION
@@ -1 +1 @@
-2.0.4
+2.0.5
diff --git a/tests/test_roleplay.py b/tests/test_roleplay.py
@@ -124,14 +124,14 @@ def test_tracking_calls(self):
  @patch("llm_roleplay.models.openai_model.OpenAIModel.generate")
  @patch("llm_roleplay.models.pipeline_model.PipelineModel.generate")
  @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.generate")
- @patch("llm_roleplay.models.openai_model.OpenAIModel._get_model")
- @patch("llm_roleplay.models.pipeline_model.PipelineModel._get_model")
- @patch("llm_roleplay.models.causal_lm_model.CausalLMModel._get_model")
+ @patch("llm_roleplay.models.openai_model.OpenAIModel.model")
+ @patch("llm_roleplay.models.pipeline_model.PipelineModel.model")
+ @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.model")
  def test_initialization(
  self,
- mock_get_model_clm,
- mock_get_model_pipe,
- mock_get_mode_openai,
+ mock_model_clm,
+ mock_model_pipe,
+ mock_model_openai,
  mock_generate_clm,
  mock_generate_pipe,
  mock_generate_openai,
@@ -185,16 +185,16 @@ def test_initialization(
  @patch("llm_roleplay.models.openai_model.OpenAIModel.generate")
  @patch("llm_roleplay.models.pipeline_model.PipelineModel.generate")
  @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.generate")
- @patch("llm_roleplay.models.openai_model.OpenAIModel._get_model")
- @patch("llm_roleplay.models.pipeline_model.PipelineModel._get_model")
- @patch("llm_roleplay.models.causal_lm_model.CausalLMModel._get_model")
+ @patch("llm_roleplay.models.openai_model.OpenAIModel.model")
+ @patch("llm_roleplay.models.pipeline_model.PipelineModel.model")
+ @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.model")
  @patch("torch.cuda.empty_cache")
  def test_resource_management(
  self,
  mock_empty_cache,
- mock_get_model_clm,
- mock_get_model_pipe,
- mock_get_mode_openai,
+ mock_model_clm,
+ mock_model_pipe,
+ mock_model_openai,
  mock_generate_clm,
  mock_generate_pipe,
  mock_generate_openai,
@@ -226,14 +226,14 @@ def test_resource_management(
  @patch("llm_roleplay.models.openai_model.OpenAIModel.generate")
  @patch("llm_roleplay.models.pipeline_model.PipelineModel.generate")
  @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.generate")
- @patch("llm_roleplay.models.openai_model.OpenAIModel._get_model")
- @patch("llm_roleplay.models.pipeline_model.PipelineModel._get_model")
- @patch("llm_roleplay.models.causal_lm_model.CausalLMModel._get_model")
+ @patch("llm_roleplay.models.openai_model.OpenAIModel.model")
+ @patch("llm_roleplay.models.pipeline_model.PipelineModel.model")
+ @patch("llm_roleplay.models.causal_lm_model.CausalLMModel.model")
  def test_dialogue_generation(
  self,
- mock_get_model_clm,
- mock_get_model_pipe,
- mock_get_mode_openai,
+ mock_model_clm,
+ mock_model_pipe,
+ mock_model_openai,
  mock_generate_clm,
  mock_generate_pipe,
  mock_generate_openai,