process_tracker_python-32 Extract Location Lookup Fix

process_tracker/process_tracker.py

@@ -157,22 +157,43 @@ def find_ready_extracts_by_filename(self, filename):
 
         return process_files
 
-    def find_ready_extracts_by_location(self, location):
+    def find_ready_extracts_by_location(self, location_name=None, location_path=None):
         """
-        For the given location name, find all matching extracts that are ready for processing
-        :param location:
-        :return:
+        For the given location path or location name, find all matching extracts that are ready for processing
+        :param location_name: The name of the location
+        :type location_name: str
+        :param location_path: The path of the location
+        :type location_path: str
+        :return: List of extract files that are in 'ready' state'.
         """
 
-        process_files = (
-            self.session.query(Extract)
-            .join(Location)
-            .join(ExtractStatus)
-            .filter(ExtractStatus.extract_status_name == "ready")
-            .filter(Location.location_name == location)
-            .order_by(Extract.extract_registration_date_time)
-            .all()
-        )
+        if location_path is not None:
+            process_files = (
+                self.session.query(Extract)
+                .join(Location)
+                .join(ExtractStatus)
+                .filter(ExtractStatus.extract_status_name == "ready")
+                .filter(Location.location_path == location_path)
+                .order_by(Extract.extract_registration_date_time)
+                .all()
+            )
+        elif location_name is not None:
+            process_files = (
+                self.session.query(Extract)
+                .join(Location)
+                .join(ExtractStatus)
+                .filter(ExtractStatus.extract_status_name == "ready")
+                .filter(Location.location_name == location_name)
+                .order_by(Extract.extract_registration_date_time)
+                .all()
+            )
+        else:
+            self.logger.error(
+                "A location name or path must be provided.  Please try again."
+            )
+            raise Exception(
+                "A location name or path must be provided.  Please try again."
+            )
 
         return process_files
 

tests/test_process_tracker.py

@@ -231,7 +231,7 @@ def test_find_ready_extracts_by_filename_partial_not_descending(self):
 
         self.assertNotEqual(expected_result, given_result)
 
-    def test_find_ready_extracts_by_location(self):
+    def test_find_ready_extracts_by_location_name(self):
         """
         Testing that for the given location name, find the extracts, provided they are in 'ready' state.  Should return
         them in ascending order by registration datettime.
@@ -266,7 +266,7 @@ def test_find_ready_extracts_by_location(self):
         ]
 
         given_result = self.process_tracker.find_ready_extracts_by_location(
-            "Test Location"
+            location_name="Test Location"
         )
         given_result = [record.full_filepath() for record in given_result]
 
@@ -307,12 +307,67 @@ def test_find_ready_extracts_by_location_not_descending(self):
         ]
 
         given_result = self.process_tracker.find_ready_extracts_by_location(
-            "Test Location"
+            location_name="Test Location"
         )
         given_result = [record.full_filepath() for record in given_result]
 
         self.assertNotEqual(expected_result, given_result)
 
+    def test_find_ready_extracts_by_location_path(self):
+        """
+        Testing that for the given location path, find the extracts, provided they are in 'ready' state.  Should return
+        them in ascending order by registration datettime.
+        :return:
+        """
+        extract = ExtractTracker(
+            process_run=self.process_tracker,
+            filename="test_extract_filename4-1.csv",
+            location_name="Test Location",
+            location_path="/home/test/extract_dir",
+        )
+
+        extract2 = ExtractTracker(
+            process_run=self.process_tracker,
+            filename="test_extract_filename4-2.csv",
+            location_name="Test Location",
+            location_path="/home/test/extract_dir",
+        )
+
+        # Need to manually change the status, because this would normally be done while the process was processing data
+        extract.extract.extract_status_id = extract.extract_status_ready
+        session = Session.object_session(extract.extract)
+        session.commit()
+
+        extract2.extract.extract_status_id = extract2.extract_status_ready
+        session = Session.object_session(extract2.extract)
+        session.commit()
+
+        expected_result = [
+            "/home/test/extract_dir/test_extract_filename4-1.csv",
+            "/home/test/extract_dir/test_extract_filename4-2.csv",
+        ]
+
+        given_result = self.process_tracker.find_ready_extracts_by_location(
+            location_path="/home/test/extract_dir"
+        )
+        given_result = [record.full_filepath() for record in given_result]
+
+        self.assertCountEqual(expected_result, given_result)
+
+    def test_find_ready_extracts_by_location_unset(self):
+        """
+        Testing that if both the location path and location name are not set, find_ready_extracts_by_location
+        will throw an exception.
+        :return:
+        """
+        with self.assertRaises(Exception) as context:
+            self.process_tracker.find_ready_extracts_by_location()
+
+        return self.assertTrue(
+            "A location name or path must be provided.  Please try again."
+            in str(context.exception)
+        )
+
     def test_find_ready_extracts_by_process(self):
         """
         Testing that for the given process name, find the extracts, provided they are in 'ready' state.