Skip to content

Commit

Permalink
Fix unit test failure
Browse files Browse the repository at this point in the history
  • Loading branch information
AnushreeBannadabhavi committed Mar 31, 2024
1 parent a4e1d55 commit f94f0bf
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions test/components/rankers/test_cohere.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,8 @@ def test_from_dict_fail_wo_env_var(self, monkeypatch):
with pytest.raises(Exception, match="The client must be instantiated be either passing in token or setting *"):
CohereRanker.from_dict(data)

def test_prepare_cohere_input_docs_default_separator(self):
def test_prepare_cohere_input_docs_default_separator(self, monkeypatch):
monkeypatch.setenv("CO_API_KEY", "fake-api-key")
component = CohereRanker(meta_fields_to_embed=["meta_field_1", "meta_field_2"])
documents = [
Document(
Expand All @@ -155,7 +156,8 @@ def test_prepare_cohere_input_docs_default_separator(self):
"meta_value_1 4\nmeta_value_2 9\ndocument number 4",
]

def test_prepare_cohere_input_docs_custom_separator(self):
def test_prepare_cohere_input_docs_custom_separator(self, monkeypatch):
monkeypatch.setenv("CO_API_KEY", "fake-api-key")
component = CohereRanker(meta_fields_to_embed=["meta_field_1", "meta_field_2"], meta_data_separator=" ")
documents = [
Document(
Expand All @@ -179,7 +181,8 @@ def test_prepare_cohere_input_docs_custom_separator(self):
"meta_value_1 4 meta_value_2 9 document number 4",
]

def test_prepare_cohere_input_docs_no_meta_data(self):
def test_prepare_cohere_input_docs_no_meta_data(self, monkeypatch):
monkeypatch.setenv("CO_API_KEY", "fake-api-key")
component = CohereRanker(meta_fields_to_embed=["meta_field_1", "meta_field_2"], meta_data_separator=" ")
documents = [Document(content=f"document number {i}") for i in range(5)]

Expand Down

0 comments on commit f94f0bf

Please sign in to comment.