Refactor tests to use fixtures (#18)

.github/workflows/pytest_workflow.yml

@@ -23,7 +23,7 @@ jobs:
  poetry config virtualenvs.in-project true
  - name: Install dependencies with Poetry
  run: |
- poetry install
+ poetry install --with api
  - name: Run tests with coverage using Poetry
  run: >
  poetry run pytest tests --cov=./ --cov-report=term

Makefile

@@ -7,14 +7,14 @@ export LANG
 
 # Initializes submodules and copies environment file sample to env file.
 .PHONY:init
-init:
- make envfile; \
- git submodule update --init; \
+init:.env
+ poetry install --with api
+ git submodule update --init;
 
 # Environment file copy
-envfile:
+.env:
 ifeq ($(wildcard envfile),)
- cp api.env.sample api.env; \
+ cp env.sample .env; \
  echo -e "\nDon't forget to update 'envfile' with all your secrets!";
 endif
 

README.md

@@ -15,8 +15,8 @@ The TA1 Middleware Service is designed to provide an intermediate job queue for
 
 ## Quickstart
 
-1. Run `make init` which will create a stub `api.env` file. 
-2. Ensure that `api.env` contains the correct endpoint and other information 
+1. Run `make init` which will create a stub `.env` file. 
+2. Ensure that `.env` contains the correct endpoint and other information 
 3. Run `make up`
 
 

api/Dockerfile

@@ -13,7 +13,7 @@ COPY --from=tmp / /
 WORKDIR /
 RUN pip install --no-cache-dir poetry==1.5.1
 RUN poetry config virtualenvs.create false && \
- poetry install --no-root --no-cache --extras api
+ poetry install --no-root --no-cache --with api
 
 COPY api ta1-service/api
 WORKDIR /ta1-service

api/server.py

@@ -4,13 +4,14 @@
 from enum import Enum
 from typing import Annotated, List, Optional
 
-from fastapi import FastAPI, HTTPException, Path, status
+from fastapi import FastAPI, HTTPException, Path, Depends, status
 from fastapi.middleware.cors import CORSMiddleware
 
 from api.models import EquationType, ExtractionJob
-from api.utils import create_job, fetch_job_status
+from api.utils import create_job, fetch_job_status, get_redis
+from lib.settings import settings
 
-LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper() # default to INFO if not set
+LOG_LEVEL = settings.LOG_LEVEL
 numeric_level = getattr(logging, LOG_LEVEL, None)
 if not isinstance(numeric_level, int):
  raise ValueError(f"Invalid log level: {LOG_LEVEL}")
@@ -43,12 +44,12 @@ def build_api(*args) -> FastAPI:
 
 
 @app.get("/status/{extraction_job_id}")
-def get_status(extraction_job_id: str) -> ExtractionJob:
+def get_status(extraction_job_id: str, redis=Depends(get_redis)) -> ExtractionJob:
  """
  Retrieve the status of a extraction
  """
 
- extraction_job_status = fetch_job_status(extraction_job_id)
+ extraction_job_status = fetch_job_status(extraction_job_id, redis)
  if (
  isinstance(extraction_job_status, int)
  and extraction_job_status == status.HTTP_404_NOT_FOUND
@@ -67,6 +68,7 @@ def equations_to_amr(
  model: str = "petrinet",
  name: Optional[str] = None,
  description: Optional[str] = None,
+ redis = Depends(get_redis)
 ) -> ExtractionJob:
  """Post equations and store an AMR to TDS
 
@@ -89,14 +91,14 @@ def equations_to_amr(
  "description": description,
  }
 
- resp = create_job(operation_name=operation_name, options=options)
+ resp = create_job(operation_name=operation_name, options=options, redis=redis)
 
  return resp
 
 
 @app.post("/code_to_amr")
 def code_to_amr(
- artifact_id: str, name: Optional[str] = None, description: Optional[str] = None
+ artifact_id: str, name: Optional[str] = None, description: Optional[str] = None, redis=Depends(get_redis)
 ) -> ExtractionJob:
  """
  Converts a code artifact to an AMR. Assumes that the code file is the first
@@ -112,13 +114,13 @@ def code_to_amr(
  operation_name = "operations.code_to_amr"
  options = {"artifact_id": artifact_id, "name": name, "description": description}
 
- resp = create_job(operation_name=operation_name, options=options)
+ resp = create_job(operation_name=operation_name, options=options, redis=redis)
 
  return resp
 
 
 @app.post("/pdf_to_text")
-def pdf_to_text(artifact_id: str) -> ExtractionJob:
+def pdf_to_text(artifact_id: str, redis=Depends(get_redis)) -> ExtractionJob:
  """Run text extractions over pdfs and stores the text as metadata on the artifact
 
  Args:
@@ -128,7 +130,7 @@ def pdf_to_text(artifact_id: str) -> ExtractionJob:
 
  options = {"artifact_id": artifact_id}
 
- resp = create_job(operation_name=operation_name, options=options)
+ resp = create_job(operation_name=operation_name, options=options, redis=redis)
 
  return resp
 
@@ -140,6 +142,7 @@ async def pdf_extractions(
  annotate_mit: bool = True,
  name: str = None,
  description: str = None,
+ redis = Depends(get_redis)
 ) -> ExtractionJob:
  """Run text extractions over pdfs
 
@@ -157,14 +160,14 @@ async def pdf_extractions(
  "description": description,
  }
 
- resp = create_job(operation_name=operation_name, options=options)
+ resp = create_job(operation_name=operation_name, options=options, redis=redis)
 
  return resp
 
 
 @app.post("/profile_dataset/{dataset_id}")
 def profile_dataset(
- dataset_id: str, artifact_id: Optional[str] = None
+ dataset_id: str, artifact_id: Optional[str] = None, redis=Depends(get_redis)
 ) -> ExtractionJob:
  """Profile dataset with MIT's profiling service. This optionally accepts an `artifact_id` which
  is expected to be some user uploaded document which has had its text extracted and stored to
@@ -183,13 +186,13 @@ def profile_dataset(
  "artifact_id": artifact_id,
  }
 
- resp = create_job(operation_name=operation_name, options=options)
+ resp = create_job(operation_name=operation_name, options=options, redis=redis)
 
  return resp
 
 
 @app.post("/profile_model/{model_id}")
-def profile_model(model_id: str, paper_artifact_id: str) -> ExtractionJob:
+def profile_model(model_id: str, paper_artifact_id: str, redis=Depends(get_redis)) -> ExtractionJob:
  """Profile model with MIT's profiling service. This takes in a paper and code artifact
  and updates a model (AMR) with the profiled metadata card. It requires that the paper
  has been extracted with `/pdf_to_text` and the code has been converted to an AMR
@@ -205,13 +208,13 @@ def profile_model(model_id: str, paper_artifact_id: str) -> ExtractionJob:
 
  options = {"model_id": model_id, "paper_artifact_id": paper_artifact_id}
 
- resp = create_job(operation_name=operation_name, options=options)
+ resp = create_job(operation_name=operation_name, options=options, redis=redis)
 
  return resp
 
 
 @app.post("/link_amr")
-def link_amr(artifact_id: str, model_id: str) -> ExtractionJob:
+def link_amr(artifact_id: str, model_id: str, redis=Depends(get_redis)) -> ExtractionJob:
  raise HTTPException(status_code=501, detail="Endpoint is under development")
 
  operation_name = "operations.link_amr"
@@ -221,6 +224,6 @@ def link_amr(artifact_id: str, model_id: str) -> ExtractionJob:
  "model_id": model_id,
  }
 
- resp = create_job(operation_name=operation_name, options=options)
+ resp = create_job(operation_name=operation_name, options=options, redis=redis)
 
  return resp

api/utils.py

@@ -17,29 +17,25 @@
 from rq.job import Job
 
 from api.models import ExtractionJob
+from lib.settings import settings
 
-LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper() # default to INFO if not set
+LOG_LEVEL = settings.LOG_LEVEL.upper()
 numeric_level = getattr(logging, LOG_LEVEL, None)
 if not isinstance(numeric_level, int):
  raise ValueError(f"Invalid log level: {LOG_LEVEL}")
 logging.basicConfig()
 logging.getLogger().setLevel(numeric_level)
 
 # REDIS CONNECTION AND QUEUE OBJECTS
-redis = Redis(
- os.environ.get("REDIS_HOST", "redis.ta1-service"),
- os.environ.get("REDIS_PORT", "6379"),
-)
-
-q = Queue(connection=redis, default_timeout=-1)
-
-
-def get_queue():
- return q
+def get_redis():
+ redis = Redis(
+ settings.REDIS_HOST,
+ settings.REDIS_PORT,
+ )
 
 
-def create_job(operation_name: str, options: Optional[Dict[Any, Any]] = None):
- q = get_queue()
+def create_job(operation_name: str, options: Optional[Dict[Any, Any]] = None, *, redis):
+ q = Queue(connection=redis, default_timeout=-1)
 
  if options is None:
  options = {}
@@ -61,7 +57,7 @@ def create_job(operation_name: str, options: Optional[Dict[Any, Any]] = None):
  if not job or force_restart:
  flattened_options = deepcopy(options)
  job = q.enqueue_call(
- func=operation_name, args=[], kwargs=flattened_options, job_id=job_id
+ func=f"worker.{operation_name}", args=[], kwargs=flattened_options, job_id=job_id
  )
  if synchronous:
  timer = 0.0
@@ -91,7 +87,7 @@ def create_job(operation_name: str, options: Optional[Dict[Any, Any]] = None):
  return ExtractionJob(id=job_id, status=status, result=result)
 
 
-def fetch_job_status(job_id):
+def fetch_job_status(job_id, redis):
  """Fetch a job's results from RQ.
 
  Args:

docker-compose.prod.yaml

@@ -13,7 +13,7 @@ services:
  ports:
  - "8010:8000"
  env_file:
- - api.env
+ - .env
  networks:
  - ta1-service
  depends_on:
@@ -35,7 +35,7 @@ services:
  context: ./
  dockerfile: workers/Dockerfile
  env_file:
- - api.env
+ - .env
  depends_on:
  - redis
  networks:

docker-compose.yaml

@@ -15,7 +15,7 @@ services:
  ports:
  - "8010:8000"
  env_file:
- - api.env
+ - .env
  networks:
  - ta1-service
  - data-api
@@ -38,7 +38,7 @@ services:
  context: ./
  dockerfile: workers/Dockerfile
  env_file:
- - api.env
+ - .env
  depends_on:
  - redis
  networks:

lib/settings.py

@@ -0,0 +1,15 @@
+from pydantic_settings import BaseSettings
+
+class Settings(BaseSettings):
+ LIVE: bool = False
+ REDIS_HOST: str = "redis.ta1-service"
+ REDIS_PORT: int = 6379
+ TA1_UNIFIED_URL: str = "http://ta1:5"
+ SKEMA_RS_URL: str = "http://skema-rs.staging.terarium.ai"
+ MIT_TR_URL: str = "http://mit:10"
+ TDS_URL: str = "http://tds:15"
+ OPENAI_API_KEY: str = "foo"
+ LOG_LEVEL: str = "INFO"
+
+
+settings = Settings()