Add basic tests.

5 files changed, 54 insertions, 5 deletions

diff --git a/.travis.yml b/.travis.yml
index 493398bc..21416ba2 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -4,7 +4,7 @@ dist: xenial
 python: "3.8"
 
 install:
-  - pip install .
+  - pip install ".[dev,test]"
 
 script:
-  - echo "Test here"
+  - pytest test
diff --git a/requirements.txt b/requirements.txt
index 8b817605..25b67425 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,6 +15,5 @@ tqdm==4.50.2
 setuptools~=50.3.2
 requests~=2.24.0
 click~=7.1.2
-bs4~=0.0.1
 beautifulsoup4~=4.9.3
 tabula-py~=2.2.0
\ No newline at end of file
diff --git a/sec_certs/download.py b/sec_certs/download.py
index 046645f9..dccbef40 100644
--- a/sec_certs/download.py
+++ b/sec_certs/download.py
@@ -22,11 +22,14 @@ def download_parallel(items: Sequence[Tuple[str, Path]], num_threads: int) -> Se
     def download(url_output):
         url, output = url_output
         return url, download_file(url, output)
+    pool = ThreadPool(num_threads)
     responses = []
     with tqdm(total=len(items)) as progress:
-        for response in ThreadPool(num_threads).imap(download, items):
+        for response in pool.imap(download, items):
             progress.update(1)
             responses.append(response)
+    pool.close()
+    pool.join()
     return responses
 
 
diff --git a/setup.py b/setup.py
index 123159b9..fc342d24 100644
--- a/setup.py
+++ b/setup.py
@@ -31,8 +31,13 @@ setup(
         "pikepdf",
         "Click",
         "requests",
-        "tqdm"
+        "tqdm",
+        "beautifulsoup4"
     ],
+    extras_require={
+        "dev": ["mypy", "flake8"],
+        "test": ["pytest", "coverage"]
+    },
     entry_points = """
         [console_scripts]
         process-certs=sec_certs.process_certificates:main
diff --git a/test/test_download.py b/test/test_download.py
new file mode 100644
index 00000000..59fef407
--- /dev/null
+++ b/test/test_download.py
@@ -0,0 +1,42 @@
+import csv
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from unittest import TestCase
+
+from sec_certs.download import download_cc_web, download_fips_web, download_cc
+
+
+class BasicTests(TestCase):
+    def setUp(self):
+        self.test_data_dir = Path(__file__).parent / "data"
+
+    def test_download(self):
+        with TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir)
+            download_cc_web(tmp_path, 4)
+            cc_files = {"cc_products_active.html", "cc_products_archived.html", "cc_labs.html",
+                        "cc_products_active.csv", "cc_products_archived.csv", "cc_pp_active.html",
+                        "cc_pp_collaborative.html", "cc_pp_archived.html", "cc_pp_active.csv",
+                        "cc_pp_archived.csv"}
+            actual = {path.name for path in tmp_path.iterdir()}
+            self.assertEqual(cc_files, actual)
+
+        with TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir)
+            download_fips_web(tmp_path)
+            fips_files = {"fips_modules_validated.html"}
+            actual = {path.name for path in tmp_path.iterdir()}
+            self.assertEqual(fips_files, actual)
+
+    def test_full_cc_download(self):
+        with open(self.test_data_dir / "certs.csv") as f:
+            reader = csv.DictReader(f)
+            certs = [(row["cert"], row["st"]) for row in reader]
+        cert_list = [( "/epfiles/" + cert, cert, "/epfiles/" + st, st) for cert, st in certs]
+        with TemporaryDirectory() as tmp_dir:
+            tmp_path = Path(tmp_dir)
+            (tmp_path / "certs").mkdir()
+            (tmp_path / "targets").mkdir()
+            download_cc(tmp_path, cert_list, 4)
+            actual = {path.name for path in tmp_path.iterdir()}
+            print(actual)