Fix GitRefs issue in a more general way

UKPLab · Jan 8, 2024 · 99d8d3b · 99d8d3b
1 parent 28f4cad
commit 99d8d3b
Showing 1 changed file with 7 additions and 6 deletions.
diff --git a/tests/test_sentence_transformer.py b/tests/test_sentence_transformer.py
@@ -83,12 +83,13 @@ def mock_list_repo_refs(self, repo_id=None, **kwargs):
             git_ref_info = GitRefInfo(name="main", ref="refs/heads/main", target_commit="123456")
         except TypeError:
             git_ref_info = GitRefInfo(dict(name="main", ref="refs/heads/main", targetCommit="123456"))
-        return GitRefs(
-            branches=[git_ref_info],
-            converts=[],
-            tags=[],
-            pull_requests=None,  # workaround for https://github.com/huggingface/huggingface_hub/issues/1956
-        )
+        # workaround for https://github.com/huggingface/huggingface_hub/issues/1956
+        git_ref_kwargs = {"branches": [git_ref_info], "converts": [], "tags": [], "pull_requests": None}
+        try:
+            return GitRefs(**git_ref_kwargs)
+        except TypeError:
+            git_ref_kwargs.pop("pull_requests")
+            return GitRefs(**git_ref_kwargs)
 
     monkeypatch.setattr(HfApi, "create_repo", mock_create_repo)
     monkeypatch.setattr(HfApi, "upload_folder", mock_upload_folder)