Add some basic tests (#8)

* Add Tests

* Refactor code to use module names when invoking external code

.github/workflows/python-checks.yaml

@@ -22,7 +22,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v3
         with:
-          python-version: '3.9'
+          python-version: '3.10'
 
       - uses: psf/black@stable
 
@@ -32,14 +32,15 @@ jobs:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v3
         with:
-          python-version: '3.9'
+          python-version: '3.10'
 
       - uses: actions/cache@v3
         with:
           path: ~/.cache/pip
           key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements-dev.txt') }}
           restore-keys: |
             ${{ runner.os }}-pip-
+
       - name: Scan
         run: |
           pip install -r requirements-dev.txt

.github/workflows/unit-tests.yaml

@@ -15,21 +15,25 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
+
       - uses: actions/cache@v3
         with:
           path: ~/.cache/pip
           key: ${{ runner.os }}-pip-${{ hashFiles('**/requirements.txt') }}
           restore-keys: |
             ${{ runner.os }}-pip-
-      - name: Set up Python 3.9
+
+      - name: Set up Python 3.10
         uses: actions/setup-python@v3
         with:
-          python-version: 3.9
+          python-version: "3.10"
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements.txt
+          pip install -r requirements-dev.txt
           pip install pytest
+
       - name: Test with Pytest unit tests
         run: |
           export DEBUG=True

gpo/api.py

@@ -0,0 +1,79 @@
+"""
+gpo rest api
+"""
+import logging
+from datetime import datetime
+import io
+import zoneinfo
+
+import fastapi
+from sqlalchemy import orm
+
+from . import crud, schemas, settings, sftp, database
+
+log = logging.getLogger(__name__)
+
+router = fastapi.APIRouter()
+
+
+def get_db():
+    """
+    get db connection
+    """
+    db = database.SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()
+
+
+@router.post("/upload", response_model=schemas.Count)
+def upload_batch(session: orm.Session = fastapi.Depends(get_db)):
+    """
+    Upload letter data file to GPO server.
+    """
+
+    letters = crud.get_letters_for_update(session)
+    count = len(letters)
+
+    if count == 0:
+        log.info("No letters in db. Nothing to upload.")
+        return {"count": count}
+
+    if settings.DEBUG:
+        output = io.StringIO()
+        sftp.write(output, letters)
+        log.debug(output.getvalue())
+        crud.delete_letters(session, letters)
+        return {"count": count}
+
+    date = datetime.now(zoneinfo.ZoneInfo("US/Eastern")).strftime("%Y%m%d")
+    file_name = f"idva-{date}-0.psv"
+
+    try:
+        sftp.write_sftp(letters, settings, file_name, settings.DEST_FILE_DIR)
+    except sftp.SftpError:
+        return fastapi.Response(status_code=400)
+
+    crud.delete_letters(session, letters)
+    log.info("Uploaded %i letter(s) as %s", count, file_name)
+
+    return {"count": count}
+
+
+@router.post("/letters", response_model=schemas.Letter)
+def queue_letter(
+    letter: schemas.LetterCreate, session: orm.Session = fastapi.Depends(get_db)
+):
+    """
+    Add a letter to the queue
+    """
+    return crud.create_letter(session, letter)
+
+
+@router.get("/letters", response_model=schemas.Count)
+def count_letter(session: orm.Session = fastapi.Depends(get_db)):
+    """
+    Get count of letter in the queue
+    """
+    return {"count": crud.count_letters(session)}

gpo/crud.py

@@ -1,36 +1,53 @@
 """
 CRUD operations for gpo µservice
 """
-from sqlalchemy import select
-from sqlalchemy.orm import Session
+import sqlalchemy
+from sqlalchemy import orm
 
 from . import models, schemas
 
 
+def count_letters(session: orm.Session) -> int:
+    """
+    count letters
+    """
+    return session.scalar(sqlalchemy.select(sqlalchemy.func.count(models.Letter.id)))
+
+
 def get_letters(
-    session: Session, skip: int = 0, limit: int = 1000
+    session: orm.Session, skip: int = 0, limit: int = 1000
 ) -> list[models.Letter]:
     """
     get letters
     """
-    return (
-        session.execute(select(models.Letter).offset(skip).limit(limit)).scalars().all()
+    statement = sqlalchemy.select(models.Letter).offset(skip).limit(limit)
+    return session.execute(statement).scalars().all()
+
+
+def get_letters_for_update(
+    session: orm.Session, skip: int = 0, limit: int = 1000
+) -> list[models.Letter]:
+    """
+    get letters with lock
+    """
+    statement = (
+        sqlalchemy.select(models.Letter).with_for_update().offset(skip).limit(limit)
     )
+    return session.execute(statement).scalars().all()
 
 
-def delete_letters(session: Session, letters: list[models.Letter]):
+def delete_letters(session: orm.Session, letters: list[models.Letter]):
     """
-    delete letters by id
+    delete list of letters by instance
     """
 
     for letter in letters:
         session.delete(letter)
 
     session.commit()
-    return
 
 
-def create_letter(session: Session, letter: schemas.LetterCreate) -> models.Letter:
+def create_letter(session: orm.Session, letter: schemas.LetterCreate) -> models.Letter:
     """
     create letter
     """

gpo/database.py

@@ -1,22 +1,18 @@
 """
 Db Connection for GPO
 """
-from sqlalchemy import create_engine
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker
-from sqlalchemy.schema import CreateSchema
-from gpo import settings
+import sqlalchemy
+from sqlalchemy import orm, schema
 
-# Sqlalchemy requires 'postgresql' as the protocol
-uri = settings.DB_URI.replace("postgres://", "postgresql://", 1)
+from . import settings
 
-schema_name = "gpo"
+engine = sqlalchemy.create_engine(
+    settings.DB_URI, connect_args={"options": f"-csearch_path={settings.SCHEMA_NAME}"}
+)
 
-engine = create_engine(uri, connect_args={"options": f"-csearch_path={schema_name}"})
+if not engine.dialect.has_schema(engine, settings.SCHEMA_NAME):
+    engine.execute(schema.CreateSchema(settings.SCHEMA_NAME))
 
-if not engine.dialect.has_schema(engine, schema_name):
-    engine.execute(CreateSchema(schema_name))
-
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine, future=True)
-
-Base = declarative_base()
+SessionLocal = orm.sessionmaker(
+    autocommit=False, autoflush=False, bind=engine, future=True
+)

gpo/main.py

@@ -1,140 +1,20 @@
 """
 GPO Microservice FastAPI Web App.
 """
-
-import datetime
-from io import StringIO
 import logging
-import math
-import csv
-from base64 import b64decode
-from zoneinfo import ZoneInfo
-from fastapi import FastAPI, Depends, Response
-import paramiko
-from starlette_prometheus import metrics, PrometheusMiddleware
-from sqlalchemy.orm import Session
-
-from . import settings, crud, models, schemas
-from .database import SessionLocal, engine
-
-# pylint: disable=invalid-name
-
-models.Base.metadata.create_all(bind=engine)
-
-app = FastAPI()
-
-app.add_middleware(PrometheusMiddleware)
-app.add_route("/metrics/", metrics)
-
-logging.getLogger().setLevel(settings.LOG_LEVEL)
-
-DEST_FILE_DIR = "gsa_order"
-
-
-def get_db():
-    """
-    get db connection
-    """
-    db = SessionLocal()
-    try:
-        yield db
-    finally:
-        db.close()
-
-
-def write(file: StringIO | paramiko.SFTPFile, letters: list[models.Letter]):
-    """
-    Write letter data to file
-
-    001|955
-    002|FAKEY MCFAKERSON|123 FAKE ST||GREAT FALLS|MT|59010|Q82GZBP71C|January 11, 2022|January 21, 2022|Example Sinatra App|https://secure.login.gov
-    003|JIMMY TESTERSON|456 FAKE RD|Apt 1|FAKE CITY|CA|40323|4WVGPG0Z5Z|January 11, 2022|January 21, 2022|Example Rails App|https://secure.login.gov
-    004|MIKE MCMOCKDATA|789 TEST AVE||FALLS CHURCH|VA|20943|5HVFT58WJ0|January 11, 2022|January 21, 2022|Example Java App|https://secure.login.gov
-    005|...
-    ...
-    956|...
-
-    """
-
-    writer = csv.writer(file, delimiter="|")
-    numLines = len(letters) + 1
-
-    # number of digits in the row index
-    numIndexDigits = math.trunc(math.log(numLines, 10)) + 1
-    # row index width is a least 2 and is enough to fit number of digits in the index
-    width = max(2, numIndexDigits)
-
-    header = [f"{1:0{width}}", len(letters)]
-    writer.writerow(header)
-
-    for i, val in enumerate(letters, start=2):
-
-        # row with index
-        row = val.as_list(f"{i:0{width}}")
-
-        # remove any pipes that might be in the data
-        sanitized_row = map(lambda x: x.replace("|", ""), row)
-        writer.writerow(sanitized_row)
-
-
-@app.post("/upload")
-def upload_batch(db: Session = Depends(get_db)):
-    """
-    Upload letter data file to GPO server.
-    """
-
-    letters = crud.get_letters(db)
-    count = len(letters)
-
-    if count == 0:
-        logging.info("No letters in db. Nothing to upload.")
-        return Response(count)
-
-    if settings.DEBUG:
-        output = StringIO()
-        write(output, letters)
-        logging.debug(output.getvalue())
-        crud.delete_letters(db, letters)
-        return Response(count)
-
-    with paramiko.SSHClient() as ssh_client:
-        host_key = paramiko.RSAKey(data=b64decode(settings.GPO_HOSTKEY))
-        ssh_client.get_host_keys().add(settings.GPO_HOST, "ssh-rsa", host_key)
-        ssh_client.connect(
-            settings.GPO_HOST,
-            username=settings.GPO_USERNAME,
-            password=settings.GPO_PASSWORD,
-        )
-        with ssh_client.open_sftp() as sftp:
-            sftp.chdir(DEST_FILE_DIR)
-            date = datetime.datetime.now(ZoneInfo("US/Eastern")).strftime("%Y%m%d")
-            try:
-                with sftp.open(f"idva-{date}-0.psv", mode="wx") as file:
-                    write(file, letters)
-            except PermissionError as err:
-                logging.error(
-                    "Error Creating file likely because it already exists: %s", err
-                )
-                return Response(status_code=500)
-
-    crud.delete_letters(db, letters)
 
-    return Response(count)
+import fastapi
+import starlette_prometheus
 
+from . import api, database, models, settings
 
-@app.post("/letters", response_model=schemas.Letter)
-def queue_letter(letter: schemas.LetterCreate, db: Session = Depends(get_db)):
-    """
-    Add a letter to the queue
-    """
+logging.basicConfig(level=settings.LOG_LEVEL)
 
-    return crud.create_letter(db, letter)
+models.Base.metadata.create_all(bind=database.engine)
 
+app = fastapi.FastAPI()
 
-@app.get("/letters")
-def count_letter(db: Session = Depends(get_db)):
-    """
-    Get count of letter in the queue
-    """
+app.add_middleware(starlette_prometheus.PrometheusMiddleware)
+app.add_route("/metrics/", starlette_prometheus.metrics)
 
-    return len(crud.get_letters(db))
+app.include_router(api.router)