mirror of
https://github.com/GSA/notifications-api.git
synced 2026-02-05 02:41:14 -05:00
merge from main
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
@@ -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",
|
||||||
|
|||||||
Reference in New Issue
Block a user