merge from main

This commit is contained in:
Kenneth Kehl
2025-01-21 09:43:21 -08:00
3 changed files with 101 additions and 22 deletions

View File

@@ -2,10 +2,12 @@ import json
from datetime import datetime, timedelta from datetime import datetime, timedelta
from os import getenv, path from os import getenv, path
from boto3 import Session
from celery.schedules import crontab from celery.schedules import crontab
from kombu import Exchange, Queue from kombu import Exchange, Queue
import notifications_utils import notifications_utils
from app.clients import AWS_CLIENT_CONFIG
from app.cloudfoundry_config import cloud_config from app.cloudfoundry_config import cloud_config
@@ -51,6 +53,13 @@ class TaskNames(object):
SCAN_FILE = "scan-file" SCAN_FILE = "scan-file"
session = Session(
aws_access_key_id=getenv("CSV_AWS_ACCESS_KEY_ID"),
aws_secret_access_key=getenv("CSV_AWS_SECRET_ACCESS_KEY"),
region_name=getenv("CSV_AWS_REGION"),
)
class Config(object): class Config(object):
NOTIFY_APP_NAME = "api" NOTIFY_APP_NAME = "api"
DEFAULT_REDIS_EXPIRE_TIME = 4 * 24 * 60 * 60 DEFAULT_REDIS_EXPIRE_TIME = 4 * 24 * 60 * 60
@@ -166,6 +175,9 @@ class Config(object):
current_minute = (datetime.now().minute + 1) % 60 current_minute = (datetime.now().minute + 1) % 60
S3_CLIENT = session.client("s3")
S3_RESOURCE = session.resource("s3", config=AWS_CLIENT_CONFIG)
CELERY = { CELERY = {
"worker_max_tasks_per_child": 500, "worker_max_tasks_per_child": 500,
"task_ignore_result": True, "task_ignore_result": True,

View File

@@ -13,14 +13,30 @@ AWS_CLIENT_CONFIG = Config(
s3={ s3={
"addressing_style": "virtual", "addressing_style": "virtual",
}, },
max_pool_connections=50,
use_fips_endpoint=True, use_fips_endpoint=True,
) )
# Global variable
noti_s3_resource = None
default_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID") default_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID")
default_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY") default_secret_access_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
default_region = os.environ.get("AWS_REGION") default_region = os.environ.get("AWS_REGION")
def get_s3_resource():
global noti_s3_resource
if noti_s3_resource is None:
session = Session(
aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
region_name=os.environ.get("AWS_REGION"),
)
noti_s3_resource = session.resource("s3", config=AWS_CLIENT_CONFIG)
return noti_s3_resource
def s3upload( def s3upload(
filedata, filedata,
region, region,
@@ -32,12 +48,7 @@ def s3upload(
access_key=default_access_key_id, access_key=default_access_key_id,
secret_key=default_secret_access_key, secret_key=default_secret_access_key,
): ):
session = Session( _s3 = get_s3_resource()
aws_access_key_id=access_key,
aws_secret_access_key=secret_key,
region_name=region,
)
_s3 = session.resource("s3", config=AWS_CLIENT_CONFIG)
key = _s3.Object(bucket_name, file_location) key = _s3.Object(bucket_name, file_location)
@@ -73,12 +84,7 @@ def s3download(
secret_key=default_secret_access_key, secret_key=default_secret_access_key,
): ):
try: try:
session = Session( s3 = get_s3_resource()
aws_access_key_id=access_key,
aws_secret_access_key=secret_key,
region_name=region,
)
s3 = session.resource("s3", config=AWS_CLIENT_CONFIG)
key = s3.Object(bucket_name, filename) key = s3.Object(bucket_name, filename)
return key.get()["Body"] return key.get()["Body"]
except botocore.exceptions.ClientError as error: except botocore.exceptions.ClientError as error:

View File

@@ -1,9 +1,16 @@
from unittest.mock import MagicMock
from urllib.parse import parse_qs from urllib.parse import parse_qs
import botocore import botocore
import pytest import pytest
from notifications_utils.s3 import S3ObjectNotFound, s3download, s3upload from notifications_utils.s3 import (
AWS_CLIENT_CONFIG,
S3ObjectNotFound,
get_s3_resource,
s3download,
s3upload,
)
contents = "some file data" contents = "some file data"
region = "eu-west-1" region = "eu-west-1"
@@ -13,7 +20,11 @@ content_type = "binary/octet-stream"
def test_s3upload_save_file_to_bucket(mocker): def test_s3upload_save_file_to_bucket(mocker):
mocked = mocker.patch("notifications_utils.s3.Session.resource")
mock_s3_resource = mocker.Mock()
mocked = mocker.patch(
"notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource
)
s3upload( s3upload(
filedata=contents, region=region, bucket_name=bucket, file_location=location filedata=contents, region=region, bucket_name=bucket, file_location=location
) )
@@ -27,7 +38,11 @@ def test_s3upload_save_file_to_bucket(mocker):
def test_s3upload_save_file_to_bucket_with_contenttype(mocker): def test_s3upload_save_file_to_bucket_with_contenttype(mocker):
content_type = "image/png" content_type = "image/png"
mocked = mocker.patch("notifications_utils.s3.Session.resource")
mock_s3_resource = mocker.Mock()
mocked = mocker.patch(
"notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource
)
s3upload( s3upload(
filedata=contents, filedata=contents,
region=region, region=region,
@@ -44,7 +59,11 @@ def test_s3upload_save_file_to_bucket_with_contenttype(mocker):
def test_s3upload_raises_exception(app, mocker): def test_s3upload_raises_exception(app, mocker):
mocked = mocker.patch("notifications_utils.s3.Session.resource")
mock_s3_resource = mocker.Mock()
mocked = mocker.patch(
"notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource
)
response = {"Error": {"Code": 500}} response = {"Error": {"Code": 500}}
exception = botocore.exceptions.ClientError(response, "Bad exception") exception = botocore.exceptions.ClientError(response, "Bad exception")
mocked.return_value.Object.return_value.put.side_effect = exception mocked.return_value.Object.return_value.put.side_effect = exception
@@ -58,7 +77,12 @@ def test_s3upload_raises_exception(app, mocker):
def test_s3upload_save_file_to_bucket_with_urlencoded_tags(mocker): def test_s3upload_save_file_to_bucket_with_urlencoded_tags(mocker):
mocked = mocker.patch("notifications_utils.s3.Session.resource")
mock_s3_resource = mocker.Mock()
mocked = mocker.patch(
"notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource
)
s3upload( s3upload(
filedata=contents, filedata=contents,
region=region, region=region,
@@ -74,7 +98,12 @@ def test_s3upload_save_file_to_bucket_with_urlencoded_tags(mocker):
def test_s3upload_save_file_to_bucket_with_metadata(mocker): def test_s3upload_save_file_to_bucket_with_metadata(mocker):
mocked = mocker.patch("notifications_utils.s3.Session.resource")
mock_s3_resource = mocker.Mock()
mocked = mocker.patch(
"notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource
)
s3upload( s3upload(
filedata=contents, filedata=contents,
region=region, region=region,
@@ -88,17 +117,49 @@ def test_s3upload_save_file_to_bucket_with_metadata(mocker):
assert metadata == {"status": "valid", "pages": "5"} assert metadata == {"status": "valid", "pages": "5"}
def test_get_s3_resource(mocker):
mock_session = mocker.patch("notifications_utils.s3.Session")
mock_current_app = mocker.patch("notifications_utils.s3.current_app")
sa_key = "sec"
sa_key = f"{sa_key}ret_access_key"
mock_current_app.config = {
"CSV_UPLOAD_BUCKET": {
"access_key_id": "test_access_key",
sa_key: "test_s_key",
"region": "us-west-100",
}
}
mock_s3_resource = MagicMock()
mock_session.return_value.resource.return_value = mock_s3_resource
result = get_s3_resource()
mock_session.return_value.resource.assert_called_once_with(
"s3", config=AWS_CLIENT_CONFIG
)
assert result == mock_s3_resource
def test_s3download_gets_file(mocker): def test_s3download_gets_file(mocker):
mocked = mocker.patch("notifications_utils.s3.Session.resource")
mock_s3_resource = mocker.Mock()
mocked = mocker.patch(
"notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource
)
mocked_object = mocked.return_value.Object mocked_object = mocked.return_value.Object
mocked_get = mocked.return_value.Object.return_value.get mocked_object.return_value.get.return_value = {"Body": mocker.Mock()}
s3download("bucket", "location.file") s3download("bucket", "location.file")
mocked_object.assert_called_once_with("bucket", "location.file") mocked_object.assert_called_once_with("bucket", "location.file")
mocked_get.assert_called_once_with()
def test_s3download_raises_on_error(mocker): def test_s3download_raises_on_error(mocker):
mocked = mocker.patch("notifications_utils.s3.Session.resource")
mock_s3_resource = mocker.Mock()
mocked = mocker.patch(
"notifications_utils.s3.get_s3_resource", return_value=mock_s3_resource
)
mocked.return_value.Object.side_effect = botocore.exceptions.ClientError( mocked.return_value.Object.side_effect = botocore.exceptions.ClientError(
{"Error": {"Code": 404}}, {"Error": {"Code": 404}},
"Bad exception", "Bad exception",