urlstore fix: compression and total_url_number()

adbar · Apr 24, 2023 · 57b5fae · 57b5fae
1 parent 3d969fa
commit 57b5fae
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/courlan/urlstore.py b/courlan/urlstore.py
@@ -392,7 +392,7 @@ def get_all_counts(self) -> List[int]:
 
     def total_url_number(self) -> int:
         "Find number of all URLs in store."
-        return sum(len(self.urldict[d].tuples) for d in self.urldict)
+        return sum(len(self._load_urls(d)) for d in self.urldict)
 
     def download_threshold_reached(self, threshold: float) -> bool:
         "Find out if the download limit (in seconds) has been reached for one of the websites in store."

diff --git a/tests/urlstore_tests.py b/tests/urlstore_tests.py
@@ -101,6 +101,7 @@ def test_urlstore():
     my_urls = UrlStore(compressed=True)
     url_buffer = UrlStore()._buffer_urls(example_urls)
     my_urls.add_urls(example_urls)
+    assert my_urls.total_url_number() == len(example_urls)
     # necessary to pickle
     my_urls._lock = None
     assert len(pickle.dumps(my_urls)) < len(pickle.dumps(url_buffer))
@@ -109,6 +110,7 @@ def test_urlstore():
     my_urls = UrlStore(compressed=True)
     url_buffer = UrlStore()._buffer_urls(test_urls)
     my_urls.add_urls(test_urls)
+    assert my_urls.total_url_number() == len(test_urls)
     # necessary to pickle
     my_urls._lock = None
     assert len(pickle.dumps(my_urls)) < len(pickle.dumps(url_buffer))