fix lint and add test

app.py

@@ -1,5 +1,4 @@
-from flask import Flask, request, jsonify
-import sys
+from flask import Flask, request
 from harvester.database.interface import HarvesterDBInterface
 from harvester.database import init_db
 from tests.database.data import new_source, new_job, new_error
@@ -13,7 +12,8 @@ def index():
     html += "<ul>"
     for rule in app.url_map.iter_rules():
         if 'static' not in rule.endpoint:
-            html += f"<li>{rule.endpoint} : <a href='{rule.rule}'>{rule.rule}</a></li><br>"
+            html += (f"<li>{rule.endpoint} : "
+                     f"<a href='{rule.rule}'>{rule.rule}</a></li><br>")
     html += "</ul>"
     return html
 

harvester/database/init_db.py

@@ -1,5 +1,5 @@
 from sqlalchemy import create_engine
-from harvester.database.models import Base, HarvestSource, HarvestJob, HarvestError
+from harvester.database.models import Base
 from sqlalchemy.engine.reflection import Inspector
 from . import DATABASE_URI
 

harvester/database/interface.py

@@ -1,11 +1,8 @@
 from sqlalchemy import create_engine, inspect
 from sqlalchemy.orm import sessionmaker
 from harvester.database.models import HarvestSource, HarvestJob, HarvestError
-import os
-from dotenv import load_dotenv
+from . import DATABASE_URI
 
-load_dotenv()
-DATABASE_URI=os.getenv("DATABASE_URI")
 engine = create_engine(DATABASE_URI)
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
@@ -27,7 +24,8 @@ def add_harvest_source(self, source_data):
 
     def get_all_harvest_sources(self):
         harvest_sources = self.db.query(HarvestSource).all()
-        harvest_sources_data = [HarvesterDBInterface._to_dict(source) for source in harvest_sources]
+        harvest_sources_data = [
+            HarvesterDBInterface._to_dict(source) for source in harvest_sources]
         return harvest_sources_data
 
     def get_harvest_source(self, source_id):
@@ -44,7 +42,8 @@ def add_harvest_job(self, job_data, source_id):
 
     def get_all_harvest_jobs(self):
         harvest_jobs = self.db.query(HarvestJob).all()
-        harvest_jobs_data = [HarvesterDBInterface._to_dict(job) for job in harvest_jobs]
+        harvest_jobs_data = [
+            HarvesterDBInterface._to_dict(job) for job in harvest_jobs]
         return harvest_jobs_data
 
     def get_harvest_job(self, job_id):
@@ -61,7 +60,8 @@ def add_harvest_error(self, error_data, job_id):
 
     def get_all_harvest_errors(self):
         harvest_errors = self.db.query(HarvestError).all()
-        harvest_errors_data = [HarvesterDBInterface._to_dict(err) for err in harvest_errors]
+        harvest_errors_data = [
+            HarvesterDBInterface._to_dict(err) for err in harvest_errors]
         return harvest_errors_data
 
     def get_harvest_error(self, error_id):

harvester/database/models.py

@@ -1,4 +1,4 @@
-from sqlalchemy import text, String, Integer, ForeignKey, ARRAY, DateTime
+from sqlalchemy import text, String, Integer, ForeignKey, DateTime
 from sqlalchemy.orm import DeclarativeBase, relationship, mapped_column
 from sqlalchemy.dialects.postgresql import JSON, UUID, ARRAY
 
@@ -20,7 +20,9 @@ class HarvestSource(Base):
 class HarvestJob(Base):
     __tablename__ = 'harvest_job'
 
-    harvest_source_id = mapped_column(UUID(as_uuid=True), ForeignKey('harvest_source.id'), nullable=False)
+    harvest_source_id = mapped_column(UUID(as_uuid=True),
+                                      ForeignKey('harvest_source.id'),
+                                      nullable=False)
     date_created = mapped_column(DateTime)
     date_finished = mapped_column(DateTime)
     records_added = mapped_column(Integer)
@@ -34,7 +36,9 @@ class HarvestJob(Base):
 class HarvestError(Base):
     __tablename__ = 'harvest_error'
 
-    harvest_job_id = mapped_column(UUID(as_uuid=True), ForeignKey('harvest_job.id'), nullable=False)
+    harvest_job_id = mapped_column(UUID(as_uuid=True),
+                                   ForeignKey('harvest_job.id'),
+                                   nullable=False)
     record_id = mapped_column(String, nullable=True)
     record_reported_id = mapped_column(String)
     date_created = mapped_column(DateTime)
@@ -44,5 +48,7 @@ class HarvestError(Base):
 
     job = relationship("HarvestJob", back_populates="errors")
 
-HarvestSource.jobs = relationship("HarvestJob", order_by=HarvestJob.id, back_populates="source")
-HarvestJob.errors = relationship("HarvestError", order_by=HarvestError.id, back_populates="job")
+HarvestSource.jobs = relationship(
+    "HarvestJob", order_by=HarvestJob.id,back_populates="source")
+HarvestJob.errors = relationship(
+    "HarvestError", order_by=HarvestError.id, back_populates="job")

tests/database/test_db.py

@@ -0,0 +1,60 @@
+import pytest
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from models import Base 
+from harvester.database.interface import HarvesterDBInterface
+
+@pytest.fixture(scope='function')
+def session():
+    engine = create_engine('sqlite:///:memory:')
+    Base.metadata.create_all(engine)
+    SessionLocal = sessionmaker(bind=engine)
+    session = SessionLocal()
+    yield session
+    session.close()
+    Base.metadata.drop_all(engine)
+
+@pytest.fixture(scope='function')
+def db_interface(session):
+    return HarvesterDBInterface(session)
+
+def test_add_harvest_source(db_interface):
+    source_data = {'name': 'Test Source', 'notification_emails': ['[email protected]']}
+    new_source = db_interface.add_harvest_source(source_data)
+    assert new_source.name == 'Test Source'
+
+def test_add_and_get_harvest_source(db_interface):
+    new_source = db_interface.add_harvest_source({
+        'name': 'Test Source',
+        'notification_emails': ['[email protected]'],
+        'organization_name': 'Test Org',
+        'frequency': 'daily',
+        'config': '{}'
+    })
+    assert new_source.name == 'Test Source' 
+
+    sources = db_interface.get_all_harvest_sources()
+    assert any(source['name'] == 'Test Source' for source in sources)
+
+
+def test_add_harvest_job(db_interface, session):
+    new_source = db_interface.add_harvest_source({
+            'name': 'Test Source',
+            'notification_emails': ['[email protected]'],
+            'organization_name': 'Test Org',
+            'frequency': 'daily',
+            'config': '{}'
+        })
+
+    job_data = {
+        'date_created': '2022-01-01',
+        'date_finished': '2022-01-02',
+        'records_added': 10,
+        'records_updated': 5,
+        'records_deleted': 2,
+        'records_errored': 1,
+        'records_ignored': 0
+    }
+    new_job = db_interface.add_harvest_job(job_data, str(new_source.id))
+    assert new_job.harvest_source_id == new_source.id
+