add file list func

ehanson8 · ehanson8 · commit 63ca6073a0e5 · 2020-01-06T13:31:39.000-05:00
diff --git a/Pipfile b/Pipfile
@@ -12,6 +12,7 @@ requests = "*"
 structlog = "*"
 attr = "*"
 click = "*"
+lxml = "*"
 
 [requires]
 python_version = "3.7"
diff --git a/Pipfile.lock b/Pipfile.lock
diff --git a/dsaps/models.py b/dsaps/models.py
@@ -1,10 +1,12 @@
 import datetime
 from functools import partial
+import os
 import operator
 import requests
 import time
 
 import attr
+from lxml import html
 import structlog
 
 op = operator.attrgetter('name')
@@ -123,6 +125,28 @@ class MetadataEntry(BaseRecord):
     language = Field()
 
 
+def build_file_list_local(directory, file_extension):
+    """Build list of files in local directory."""
+    file_list = {}
+    for root, dirs, files in os.walk(directory, topdown=True):
+        for file in files:
+            if file.endswith(file_extension):
+                full_file_path = os.path.join(root, file)
+                file_list[file] = full_file_path
+    return file_list
+
+
+def build_file_list_remote(directory_url, file_extension):
+    """Build list of files in local directory."""
+    file_list = {}
+    response = requests.get(directory_url)
+    links = html.fromstring(response.content).iterlinks()
+    for link in links:
+        if link[2].endswith(file_extension):
+            file_list[link[2]] = f'{directory_url}{link[2]}'
+    return file_list
+
+
 def elapsed_time(start_time, label):
     """Calculate elapsed time."""
     td = datetime.timedelta(seconds=time.time() - start_time)
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -1,3 +1,5 @@
+import os
+
 import attr
 import pytest
 import requests_mock
@@ -60,6 +62,7 @@ def test_filtered_item_search(client):
 
 
 def test__pop_inst(client):
+    """Test _pop_inst function."""
     class_type = models.Collection
     rec_obj = {'name': 'Test title', 'type': 'collection', 'items': []}
     rec_obj = client._pop_inst(class_type, rec_obj)
@@ -68,7 +71,36 @@ def test__pop_inst(client):
 
 
 def test__build_uuid_list(client):
+    """Test _build_uuid_list function."""
     rec_obj = {'items': [{'uuid': '1234'}]}
     children = 'items'
     child_list = client._build_uuid_list(rec_obj, children)
     assert '1234' in child_list
+
+
+def test_build_file_list_local():
+    """Test filtered_item_search function."""
+    file_extension = 'pdf'
+    directory = 'test_temp'
+    os.mkdir(directory)
+    open(f'{directory}/999.pdf', 'w')
+    file_list = models.build_file_list_local(directory, file_extension)
+    os.remove(f'{directory}/999.pdf')
+    os.rmdir(directory)
+    assert '999.pdf' in file_list
+
+
+def test_build_file_list_remote():
+    """Test build_file_list_remote function."""
+    content = '<!DOCTYPE HTML PUBLIC "-//W3C//DTD HTML 3.2 Final//EN"><html>'
+    content += '<head><title>Index of /pdf</title></head><body><h1>Index of /'
+    content += 'pdf</h1><table><tr><th>Name</th><th>Last modified</th><th>'
+    content += 'Size</th></tr><tr><td><a href="999.pdf">999.pdf</a></td><td>'
+    content += '2001-02-16 11:59 </td><td>107K</td></tr></table></body></html>'
+    with requests_mock.Mocker() as m:
+        directory_url = 'mock://test.com/pdfs/'
+        file_extension = 'pdf'
+        m.get(directory_url, text=content)
+        file_list = models.build_file_list_remote(directory_url,
+                                                  file_extension)
+        assert '999.pdf' in file_list
