fix test

janeklb · Aug 3, 2024 · ce3b23f · ce3b23f
1 parent dbfb224
commit ce3b23f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 17 deletions.
diff --git a/ghsearch/gh_search.py b/ghsearch/gh_search.py
@@ -40,16 +40,11 @@ def _confirm_continue_near_limit(core_rate: Rate, num_results: int, calls_per_re
     )
 
 
-def _echo_rate_limits(rate_limit: RateLimit | None) -> None:
-    if rate_limit is None:
-        click.echo("(Rate limiting is disabled)")
-    else:
-        core = rate_limit.core
-        search = rate_limit.search
-        click.echo(
-            f"Core rate limit: {core.remaining}/{core.limit} (resets {core.reset}), "
-            f"Search rate limit: {search.remaining}/{search.limit} (resets {search.reset})"
-        )
+def _echo_rate_limits(rate_limit: RateLimit) -> None:
+    click.echo(
+        f"Core rate limit: {rate_limit.core.remaining}/{rate_limit.core.limit} (resets {rate_limit.core.reset}), "
+        f"Search rate limit: {rate_limit.search.remaining}/{rate_limit.search.limit} (resets {rate_limit.search.reset})"
+    )
 
 
 class GHSearch:
@@ -70,7 +65,7 @@ def get_rate_limit(self) -> RateLimit | None:
     def get_filtered_results(self, query: List[str]) -> List[ContentFile]:
         rate_limit = self.get_rate_limit()
 
-        if self.verbose:
+        if rate_limit and self.verbose:
             _echo_rate_limits(rate_limit)
 
         results = self.client.search_code(query=" ".join(query))
@@ -92,8 +87,8 @@ def get_filtered_results(self, query: List[str]) -> List[ContentFile]:
                     elif self.verbose:
                         click.echo(f"Skipping result for {result.repository.full_name} via {exclude_reason}")
 
-        if self.verbose:
-            rate_limit = self.get_rate_limit()
+        rate_limit = self.get_rate_limit()
+        if rate_limit and self.verbose:
             _echo_rate_limits(rate_limit)
 
         return filtered_results

diff --git a/tests/unit/gh_search_test.py b/tests/unit/gh_search_test.py
@@ -129,13 +129,13 @@ def test_get_filtered_results_many_calls(mock_client, mock_click):
     )
 
 
-def test_get_filtered_results_rate_limiting_disabled(mock_client, mock_click):
+def test_get_filtered_results_rate_limiting_disabled(mock_client):
     mock_client.get_rate_limit.side_effect = github.GithubException(404, "Not Found")
-    mock_client.search_code.return_value = MockPaginatedList(*[], total_count=0)
     mock_filter = Mock()
     mock_filter.uses_core_api = True
 
-    ghsearch = GHSearch(mock_client, [mock_filter], True)
+    ghsearch = GHSearch(mock_client, [])
     ghsearch.get_filtered_results(["query", "org:bort"])
 
-    mock_click.echo.assert_any_call("(Rate limiting is disabled)")
+    # ensure get_rate_limit was called (and the side_effect above handled)
+    mock_client.get_rate_limit.assert_called()