From 78d751afb0604a7e4ff235cfb3dd07c3965deda5 Mon Sep 17 00:00:00 2001 From: Brianna Smart Date: Fri, 13 Dec 2024 11:54:57 -0800 Subject: [PATCH] Update tests and remove more gcs references --- alertdb/bin/alertdb.py | 8 ++++---- alertdb/storage.py | 7 +++---- tests/test_integration.py | 8 ++++---- 3 files changed, 11 insertions(+), 12 deletions(-) diff --git a/alertdb/bin/alertdb.py b/alertdb/bin/alertdb.py index 2aa5d66..12e8232 100644 --- a/alertdb/bin/alertdb.py +++ b/alertdb/bin/alertdb.py @@ -4,7 +4,7 @@ import uvicorn from alertdb.server import create_server -from alertdb.storage import FileBackend, GoogleObjectStorageBackend +from alertdb.storage import FileBackend, USDFObjectStorageBackend logger = logging.getLogger(__name__) @@ -50,13 +50,13 @@ def main(): "--gcp-bucket-alerts", type=str, default="alert-packets", - help="when using the google-cloud backend, the name of the GCS bucket for alert packets", + help="when using the google-cloud backend, the name of the boto3 bucket for alert packets", ) parser.add_argument( "--gcp-bucket-schemas", type=str, default="alert-schemas", - help="when using the google-cloud backend, the name of the GCS bucket for alert schemas", + help="when using the google-cloud backend, the name of the boto3 bucket for alert schemas", ) parser.add_argument("--verbose", action="store_true", help="log a bunch") parser.add_argument("--debug", action="store_true", help="log even more") @@ -99,7 +99,7 @@ def main(): logger.info("gcp_bucket_alerts: %s", args.gcp_bucket_alerts) logger.info("gcp_bucket_schemas: %s", args.gcp_bucket_schemas) - backend = GoogleObjectStorageBackend( + backend = USDFObjectStorageBackend( args.gcp_project, args.gcp_bucket_alerts, args.gcp_bucket_schemas ) diff --git a/alertdb/storage.py b/alertdb/storage.py index fc2a884..92ccb64 100644 --- a/alertdb/storage.py +++ b/alertdb/storage.py @@ -7,9 +7,6 @@ import os.path import boto3 -import botocore -import google.api_core.exceptions -import google.cloud.storage as gcs logger = logging.getLogger(__name__) @@ -147,7 +144,9 @@ def __init__( def get_alert(self, alert_id: str) -> bytes: logger.info("retrieving alert id=%s", alert_id) try: - alert_key = f"/alert_archive/v1/alerts/{alert_id}.avro.gz" # boto3 terminology for objects, objects live in prefixes inside of buckets + alert_key = f"/alert_archive/v1/alerts/{alert_id}.avro.gz" + # boto3 terminology for objects, objects live in prefixes inside + # of buckets blob = self.object_store_client.get_object( Bucket=self.packet_bucket, Key=alert_key ) diff --git a/tests/test_integration.py b/tests/test_integration.py index 727a3d7..b3713cb 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -2,12 +2,12 @@ import os import unittest -import google.cloud.storage as gcs +import boto3 import requests from fastapi.testclient import TestClient from alertdb.server import create_server -from alertdb.storage import GoogleObjectStorageBackend +from alertdb.storage import USDFObjectStorageBackend logger = logging.getLogger(__name__) logger.level = logging.DEBUG @@ -29,7 +29,7 @@ def setUpClass(cls): ) packet_bucket_name = "alertdb_server_integration_test_bucket_packets" schema_bucket_name = "alertdb_server_integration_test_bucket_schemas" - client = gcs.Client(project=gcp_project) + client = boto3.client("s3", endpoint_url=gcp_project) logger.info("creating bucket %s", packet_bucket_name) packet_bucket = client.create_bucket(packet_bucket_name) @@ -89,7 +89,7 @@ def setUp(self): """ Run a local instance of the server. """ - backend = GoogleObjectStorageBackend( + backend = USDFObjectStorageBackend( self.gcp_project, self.packet_bucket_name, self.schema_bucket_name ) self.server = create_server(backend)