diff --git a/tests/api.py b/tests/api.py index 2708ad8..1c90403 100644 --- a/tests/api.py +++ b/tests/api.py @@ -1,3 +1,5 @@ +import pathlib +import tempfile import unittest from unittest.mock import patch, MagicMock from yeti.api import YetiApi @@ -295,22 +297,27 @@ def test_download_dfiq_archive(self, mock_post): json={"count": 0}, ) - @patch("yeti.api.requests.Session.post") - def test_upload_dfiq_archive(self, mock_post): + # Mocks Session.send rather than Session.post, so that requests builds the + # request and rejects arguments it doesn't accept. The archive is a real + # file because requests may open files of its own, such as ~/.netrc. + @patch("yeti.api.requests.Session.send") + def test_upload_dfiq_archive(self, mock_send): mock_response = MagicMock() mock_response.content = b'{"uploaded": 1}' - mock_post.return_value = mock_response - - with patch("builtins.open", unittest.mock.mock_open(read_data=b"data")): - result = self.api.upload_dfiq_archive("path/to/archive.zip") - self.assertEqual(result, {"uploaded": 1}) - self.assertEqual( - mock_post.call_args[0][0], "http://fake-url/api/v2/dfiq/from_archive" - ) + mock_send.return_value = mock_response + + with tempfile.TemporaryDirectory() as tmp_dir: + archive_path = pathlib.Path(tmp_dir) / "archive.zip" + archive_path.write_bytes(b"zip-archive-bytes") + result = self.api.upload_dfiq_archive(str(archive_path)) + self.assertEqual(result, {"uploaded": 1}) + request = mock_send.call_args[0][0] + self.assertEqual(request.url, "http://fake-url/api/v2/dfiq/from_archive") self.assertRegex( - mock_post.call_args[1]["headers"]["Content-Type"], + request.headers["Content-Type"], "multipart/form-data; boundary=[a-f0-9]{32}", ) + self.assertIn(b"zip-archive-bytes", request.body) @patch("yeti.api.requests.Session.post") def test_add_observable(self, mock_post): diff --git a/tests/e2e.py b/tests/e2e.py index 4409fde..b08d398 100644 --- a/tests/e2e.py +++ b/tests/e2e.py @@ -1,7 +1,10 @@ import os +import pathlib +import tempfile import time import unittest import uuid +import zipfile from yeti import errors from yeti.api import YetiApi @@ -237,6 +240,30 @@ def test_dfiq_from_yaml_and_patch(self): patched = self.api.patch_dfiq(patched) self.assertEqual(patched["description"], "patched as an object") + def test_upload_dfiq_archive(self): + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + scenario_uuid = str(uuid.uuid4()) + scenario_yaml = f"""--- +name: testArchiveScenario +type: scenario +description: > + test +id: S1991 +uuid: {scenario_uuid} +dfiq_version: 1.1.0 +""" + with tempfile.TemporaryDirectory() as tmp_dir: + archive_path = pathlib.Path(tmp_dir) / "dfiq.zip" + # The server reads YAML files one directory deep in the archive. + with zipfile.ZipFile(archive_path, "w") as archive: + archive.writestr("scenarios/S1991.yaml", scenario_yaml) + result = self.api.upload_dfiq_archive(str(archive_path)) + + self.assertEqual(result, {"total_added": 1}) + scenario = self.api.find_dfiq(name="testArchiveScenario", dfiq_type="scenario") + self.assertIsNotNone(scenario) + self.assertEqual(scenario["uuid"], scenario_uuid) + def test_new_tag(self): self.api.auth_api_key(os.getenv("YETI_API_KEY")) tag = self.api.new_tag("testTag", description="test") diff --git a/yeti/api.py b/yeti/api.py index 12b5974..f98cdd6 100644 --- a/yeti/api.py +++ b/yeti/api.py @@ -133,7 +133,7 @@ def do_request( if json_data: request_kwargs["json"] = json_data if body: - request_kwargs["body"] = body + request_kwargs["data"] = body if params: url = f"{url}?{urllib.parse.urlencode(params)}"