Skip to content

Commit

Permalink
fix gcs tus writer (#2392)
Browse files Browse the repository at this point in the history
  • Loading branch information
drf7 authored Aug 9, 2024
1 parent c10226c commit adb1515
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 7 deletions.
13 changes: 10 additions & 3 deletions nucliadb/src/nucliadb/writer/tus/gcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@
import asyncio
import base64
import json
import os
import socket
import tempfile
import uuid
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
Expand All @@ -34,7 +36,7 @@
import google.auth.compute_engine.credentials # type: ignore
import google.auth.transport.requests # type: ignore
from google.auth.exceptions import DefaultCredentialsError # type: ignore
from google.oauth2 import service_account
from oauth2client.service_account import ServiceAccountCredentials # type: ignore

from nucliadb.writer import logger
from nucliadb.writer.tus.dm import FileDataManager
Expand Down Expand Up @@ -123,8 +125,13 @@ async def initialize(
self._credentials = None

if self.json_credentials is not None and self.json_credentials.strip() != "":
self._credentials = service_account.Credentials.from_service_account_info(
base64.b64decode(self.json_credentials).decode("utf-8"), scopes=SCOPES
self.json_credentials_file = os.path.join(tempfile.mkdtemp(), "gcs_credentials.json")
with open(self.json_credentials_file, "w") as file:
file.write(
base64.b64decode(self.json_credentials).decode("utf-8")
)
self._credentials = ServiceAccountCredentials.from_json_keyfile_name(
self.json_credentials_file, SCOPES
)
else:
try:
Expand Down
8 changes: 4 additions & 4 deletions nucliadb/tests/writer/unit/tus/test_gcs.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

from unittest.mock import MagicMock, patch
from unittest.mock import ANY, MagicMock, patch

import pytest

Expand All @@ -26,7 +26,7 @@

@pytest.mark.asyncio
@patch("nucliadb.writer.tus.gcs.aiohttp")
@patch("nucliadb.writer.tus.gcs.service_account")
@patch("nucliadb.writer.tus.gcs.ServiceAccountCredentials")
async def test_tus_gcs(mock_sa, mock_aiohttp):
mock_session = MagicMock()
mock_aiohttp.ClientSession.return_value = mock_session
Expand All @@ -44,8 +44,8 @@ async def test_tus_gcs(mock_sa, mock_aiohttp):
json_credentials="dGVzdC1jcmVk",
)

mock_sa.Credentials.from_service_account_info.assert_called_once_with(
"test-cred", scopes=["https://www.googleapis.com/auth/devstorage.read_write"]
mock_sa.from_json_keyfile_name.assert_called_once_with(
ANY, ["https://www.googleapis.com/auth/devstorage.read_write"]
)

assert await gblobstore.check_exists("test-bucket")
Expand Down

0 comments on commit adb1515

Please sign in to comment.