Wrote tests for ProcessTracking find extracts functions

✅ Wrote tests for ProcessTracking

Wrote tests for helper functions that are used by ProcessTracking to
find extracts by process name, location name, or full/partial filename.

They are not necessarily part of the core functionality of the ProcessTracking class, but by having them as part of it, users will have an easier time using the framework.

Author: OpenDataAlex

models/extract.py

@@ -34,6 +34,7 @@ class Extract(Base):
 #    extract_archive_date_time = Column(DateTime, nullable=False, default=default_date)
 
     extract_process = relationship("ExtractProcess", back_populates='process_extracts')
+    locations = relationship("Location", foreign_keys=[extract_location_id])
 
 
 class ExtractProcess(Base):

process_tracker/extract_tracking.py

@@ -58,6 +58,9 @@ def __init__(self, process_run, filename, location_path, location_name=None):
 
         self.extract_process = self.retrieve_extract_process()
 
+        self.extract.extract_status_id = self.extract_status_initializing
+        session.commit()
+
     def change_extract_status(self, new_status):
         """
         Change an extract record status.

process_tracker/process_tracking.py

@@ -9,7 +9,7 @@
 from process_tracker.data_store import DataStore
 
 from models.actor import Actor
-from models.extract import Extract, ExtractStatus, Location
+from models.extract import Extract, ExtractProcess, ExtractStatus, Location
 from models.process import ErrorTracking, ErrorType, Process, ProcessTracking, ProcessStatus, ProcessType
 from models.source import Source
 from models.tool import Tool
@@ -97,14 +97,14 @@ def find_ready_extracts_by_filename(filename):
         """
         extract_files = []
 
-        process_files = session.query(Extract, Location)\
+        process_files = session.query(Extract.extract_filename, Location.location_path)\
                                .join(Location)\
                                .join(ExtractStatus)\
                                .filter(Extract.extract_filename.like("%" + filename + "%"))\
                                .filter(ExtractStatus.extract_status_name == 'ready')
 
         for record in process_files:
-            extract_files.append(join(record.location_path, record.filename))
+            extract_files.append(join(record.location_path, record.extract_filename))
 
         return extract_files
 
@@ -117,35 +117,36 @@ def find_ready_extracts_by_location(location):
         """
         extract_files = []
 
-        process_files = session.query(Extract, Location)\
+        process_files = session.query(Extract.extract_filename, Location.location_path)\
                                .join(Location)\
                                .join(ExtractStatus)\
                                .filter(ExtractStatus.extract_status_name == 'ready')\
                                .filter(Location.location_name == location)
 
         for record in process_files:
-            extract_files.append(join(record.location_path, record.filename))
+            extract_files.append(join(record.location_path, record.extract_filename))
 
         return extract_files
 
     @staticmethod
-    def find_ready_extracts_by_process(self, extract_process_name):
+    def find_ready_extracts_by_process(extract_process_name):
         """
         For the given named process, find the extracts that are ready for processing.
         :return: List of OS specific filepaths with filenames.
         """
         extract_files = []
 
-        process_files = session.query(Extract, Location) \
-            .join(Location) \
+        process_files = session.query(Extract.extract_filename, Location.location_path) \
+            .join(ExtractStatus, Extract.extract_status_id == ExtractStatus.extract_status_id) \
+            .join(Location, Extract.extract_location_id == Location.location_id) \
+            .join(ExtractProcess, Extract.extract_id == ExtractProcess.extract_tracking_id) \
             .join(ProcessTracking) \
             .join(Process) \
-            .join(ExtractStatus)\
             .filter(Process.process_name == extract_process_name
                     , ExtractStatus.extract_status_name == 'ready')
 
         for record in process_files:
-            extract_files.append(join(record.location_path, record.filename))
+            extract_files.append(join(record.location_path, record.extract_filename))
 
         return extract_files
 

tests/test_process_tracking.py

@@ -3,10 +3,11 @@
 from datetime import datetime
 import unittest
 
+from models.extract import Extract, ExtractProcess
 from models.process import ErrorType, ErrorTracking, Process, ProcessTracking
 
-from process_tracker import data_store_type
-from process_tracker import session
+from process_tracker import data_store_type, session
+from process_tracker.extract_tracking import ExtractTracker
 from process_tracker.process_tracking import ProcessTracker
 
 
@@ -36,7 +37,9 @@ def tearDown(self):
         Need to clean up tables to return them to pristine state for other tests.
         :return:
         """
+        session.query(ExtractProcess).delete()
         session.query(ProcessTracking).delete()
+        session.query(Extract).delete()
         session.query(ErrorTracking).delete()
         session.query(ErrorType).delete()
         session.commit()
@@ -52,6 +55,90 @@ def test_verify_session_variables_postgresql(self):
 
         self.assertEqual(expected_result, given_result)
 
+    def test_find_ready_extracts_by_filename_full(self):
+        """
+        Testing that for the given full filename, find the extract, provided it's in 'ready' state.
+        :return:
+        """
+        extract = ExtractTracker(process_run=self.process_tracker
+                       , filename='test_extract_filename2.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.commit()
+
+        expected_result = ['/home/test/extract_dir/test_extract_filename2.csv']
+
+        given_result = self.process_tracker.find_ready_extracts_by_filename('test_extract_filename2.csv')
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_find_ready_extracts_by_filename_partial(self):
+        """
+        Testing that for the given partial filename, find the extracts, provided they are in 'ready' state.
+        :return:
+        """
+        extract = ExtractTracker(process_run=self.process_tracker
+                       , filename='test_extract_filename3.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.commit()
+
+        expected_result = ['/home/test/extract_dir/test_extract_filename3.csv']
+
+        given_result = self.process_tracker.find_ready_extracts_by_filename('test_extract_filename')
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_find_ready_extracts_by_location(self):
+        """
+        Testing that for the given location name, find the extracts, provided they are in 'ready' state.
+        :return:
+        """
+        extract = ExtractTracker(process_run=self.process_tracker
+                       , filename='test_extract_filename4.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.commit()
+
+        expected_result = ['/home/test/extract_dir/test_extract_filename4.csv']
+
+        given_result = self.process_tracker.find_ready_extracts_by_location('Test Location')
+
+        self.assertEqual(expected_result, given_result)
+
+    def test_find_ready_extracts_by_process(self):
+        """
+        Testing that for the given process name, find the extracts, provided they are in 'ready' state.
+        :return:
+        """
+        extract = ExtractTracker(process_run=self.process_tracker
+                       , filename='test_extract_filename5.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.commit()
+
+        expected_result = ['/home/test/extract_dir/test_extract_filename5.csv']
+
+        given_result = self.process_tracker.find_ready_extracts_by_process('Testing Process Tracking Initialization')
+
+        self.assertEqual(expected_result, given_result)
+
     def test_initializing_process_tracking(self):
         """
         Testing that when ProcessTracking is initialized, the necessary objects are created.