mirror of
https://github.com/zulip/zulip.git
synced 2025-11-19 05:58:25 +00:00
dcbb2a78ca
JsonableError has two major benefits over json_error: * It can be raised from anywhere in the codebase, rather than being a return value, which is much more convenient for refactoring, as one doesn't potentially need to change error handling style when extracting a bit of view code to a function. * It is guaranteed to contain the `code` property, which is helpful for API consistency. Various stragglers are not updated because JsonableError requires subclassing in order to specify custom data or HTTP status codes.
221 lines
8.8 KiB
Python
221 lines
8.8 KiB
Python
import os
|
|
from unittest.mock import patch
|
|
|
|
import botocore.exceptions
|
|
import orjson
|
|
from django.conf import settings
|
|
from django.utils.timezone import now as timezone_now
|
|
|
|
from analytics.models import RealmCount
|
|
from zerver.lib.exceptions import JsonableError
|
|
from zerver.lib.test_classes import ZulipTestCase
|
|
from zerver.lib.test_helpers import (
|
|
create_dummy_file,
|
|
create_s3_buckets,
|
|
stdout_suppressed,
|
|
use_s3_backend,
|
|
)
|
|
from zerver.models import RealmAuditLog
|
|
from zerver.views.realm_export import export_realm
|
|
|
|
|
|
class RealmExportTest(ZulipTestCase):
|
|
"""
|
|
API endpoint testing covers the full end-to-end flow
|
|
from both the S3 and local uploads perspective.
|
|
|
|
`test_endpoint_s3` and `test_endpoint_local_uploads` follow
|
|
an identical pattern, which is documented in both test
|
|
functions.
|
|
"""
|
|
|
|
def test_export_as_not_admin(self) -> None:
|
|
user = self.example_user("hamlet")
|
|
self.login_user(user)
|
|
with self.assertRaises(JsonableError):
|
|
export_realm(self.client_post, user)
|
|
|
|
@use_s3_backend
|
|
def test_endpoint_s3(self) -> None:
|
|
admin = self.example_user("iago")
|
|
self.login_user(admin)
|
|
bucket = create_s3_buckets(settings.S3_AVATAR_BUCKET)[0]
|
|
tarball_path = create_dummy_file("test-export.tar.gz")
|
|
|
|
# Test the export logic.
|
|
with patch("zerver.lib.export.do_export_realm", return_value=tarball_path) as mock_export:
|
|
with self.settings(LOCAL_UPLOADS_DIR=None), stdout_suppressed(), self.assertLogs(
|
|
level="INFO"
|
|
) as info_logs:
|
|
result = self.client_post("/json/export/realm")
|
|
self.assertTrue("INFO:root:Completed data export for zulip in " in info_logs.output[0])
|
|
self.assert_json_success(result)
|
|
self.assertFalse(os.path.exists(tarball_path))
|
|
args = mock_export.call_args_list[0][1]
|
|
self.assertEqual(args["realm"], admin.realm)
|
|
self.assertEqual(args["public_only"], True)
|
|
self.assertIn("/tmp/zulip-export-", args["output_dir"])
|
|
self.assertEqual(args["threads"], 6)
|
|
|
|
# Get the entry and test that iago initiated it.
|
|
audit_log_entry = RealmAuditLog.objects.filter(
|
|
event_type=RealmAuditLog.REALM_EXPORTED
|
|
).first()
|
|
self.assertEqual(audit_log_entry.acting_user_id, admin.id)
|
|
|
|
# Test that the file is hosted, and the contents are as expected.
|
|
export_path = orjson.loads(audit_log_entry.extra_data)["export_path"]
|
|
assert export_path.startswith("/")
|
|
path_id = export_path[1:]
|
|
self.assertEqual(bucket.Object(path_id).get()["Body"].read(), b"zulip!")
|
|
|
|
result = self.client_get("/json/export/realm")
|
|
self.assert_json_success(result)
|
|
|
|
# Test that the export we have is the export we created.
|
|
export_dict = result.json()["exports"]
|
|
self.assertEqual(export_dict[0]["id"], audit_log_entry.id)
|
|
self.assertEqual(
|
|
export_dict[0]["export_url"],
|
|
"https://test-avatar-bucket.s3.amazonaws.com" + export_path,
|
|
)
|
|
self.assertEqual(export_dict[0]["acting_user_id"], admin.id)
|
|
self.assert_length(
|
|
export_dict,
|
|
RealmAuditLog.objects.filter(
|
|
realm=admin.realm, event_type=RealmAuditLog.REALM_EXPORTED
|
|
).count(),
|
|
)
|
|
|
|
# Finally, delete the file.
|
|
result = self.client_delete(f"/json/export/realm/{audit_log_entry.id}")
|
|
self.assert_json_success(result)
|
|
with self.assertRaises(botocore.exceptions.ClientError):
|
|
bucket.Object(path_id).load()
|
|
|
|
# Try to delete an export with a `deleted_timestamp` key.
|
|
audit_log_entry.refresh_from_db()
|
|
export_data = orjson.loads(audit_log_entry.extra_data)
|
|
self.assertIn("deleted_timestamp", export_data)
|
|
result = self.client_delete(f"/json/export/realm/{audit_log_entry.id}")
|
|
self.assert_json_error(result, "Export already deleted")
|
|
|
|
# Now try to delete a non-existent export.
|
|
result = self.client_delete("/json/export/realm/0")
|
|
self.assert_json_error(result, "Invalid data export ID")
|
|
|
|
def test_endpoint_local_uploads(self) -> None:
|
|
admin = self.example_user("iago")
|
|
self.login_user(admin)
|
|
tarball_path = create_dummy_file("test-export.tar.gz")
|
|
|
|
# Test the export logic.
|
|
with patch("zerver.lib.export.do_export_realm", return_value=tarball_path) as mock_export:
|
|
with stdout_suppressed(), self.assertLogs(level="INFO") as info_logs:
|
|
result = self.client_post("/json/export/realm")
|
|
self.assertTrue("INFO:root:Completed data export for zulip in " in info_logs.output[0])
|
|
self.assert_json_success(result)
|
|
self.assertFalse(os.path.exists(tarball_path))
|
|
args = mock_export.call_args_list[0][1]
|
|
self.assertEqual(args["realm"], admin.realm)
|
|
self.assertEqual(args["public_only"], True)
|
|
self.assertIn("/tmp/zulip-export-", args["output_dir"])
|
|
self.assertEqual(args["threads"], 6)
|
|
|
|
# Get the entry and test that iago initiated it.
|
|
audit_log_entry = RealmAuditLog.objects.filter(
|
|
event_type=RealmAuditLog.REALM_EXPORTED
|
|
).first()
|
|
self.assertEqual(audit_log_entry.acting_user_id, admin.id)
|
|
|
|
# Test that the file is hosted, and the contents are as expected.
|
|
export_path = orjson.loads(audit_log_entry.extra_data).get("export_path")
|
|
response = self.client_get(export_path)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assert_url_serves_contents_of_file(export_path, b"zulip!")
|
|
|
|
result = self.client_get("/json/export/realm")
|
|
self.assert_json_success(result)
|
|
|
|
# Test that the export we have is the export we created.
|
|
export_dict = result.json()["exports"]
|
|
self.assertEqual(export_dict[0]["id"], audit_log_entry.id)
|
|
self.assertEqual(export_dict[0]["export_url"], admin.realm.uri + export_path)
|
|
self.assertEqual(export_dict[0]["acting_user_id"], admin.id)
|
|
self.assert_length(
|
|
export_dict,
|
|
RealmAuditLog.objects.filter(
|
|
realm=admin.realm, event_type=RealmAuditLog.REALM_EXPORTED
|
|
).count(),
|
|
)
|
|
|
|
# Finally, delete the file.
|
|
result = self.client_delete(f"/json/export/realm/{audit_log_entry.id}")
|
|
self.assert_json_success(result)
|
|
response = self.client_get(export_path)
|
|
self.assertEqual(response.status_code, 404)
|
|
|
|
# Try to delete an export with a `deleted_timestamp` key.
|
|
audit_log_entry.refresh_from_db()
|
|
export_data = orjson.loads(audit_log_entry.extra_data)
|
|
self.assertIn("deleted_timestamp", export_data)
|
|
result = self.client_delete(f"/json/export/realm/{audit_log_entry.id}")
|
|
self.assert_json_error(result, "Export already deleted")
|
|
|
|
# Now try to delete a non-existent export.
|
|
result = self.client_delete("/json/export/realm/0")
|
|
self.assert_json_error(result, "Invalid data export ID")
|
|
|
|
def test_realm_export_rate_limited(self) -> None:
|
|
admin = self.example_user("iago")
|
|
self.login_user(admin)
|
|
|
|
current_log = RealmAuditLog.objects.filter(event_type=RealmAuditLog.REALM_EXPORTED)
|
|
self.assert_length(current_log, 0)
|
|
|
|
exports = []
|
|
for i in range(0, 5):
|
|
exports.append(
|
|
RealmAuditLog(
|
|
realm=admin.realm,
|
|
event_type=RealmAuditLog.REALM_EXPORTED,
|
|
event_time=timezone_now(),
|
|
)
|
|
)
|
|
RealmAuditLog.objects.bulk_create(exports)
|
|
|
|
with self.assertRaises(JsonableError) as error:
|
|
export_realm(self.client_post, admin)
|
|
self.assertEqual(str(error.exception), "Exceeded rate limit.")
|
|
|
|
def test_upload_and_message_limit(self) -> None:
|
|
admin = self.example_user("iago")
|
|
self.login_user(admin)
|
|
realm_count = RealmCount.objects.create(
|
|
realm_id=admin.realm.id,
|
|
end_time=timezone_now(),
|
|
subgroup=1,
|
|
value=0,
|
|
property="messages_sent:client:day",
|
|
)
|
|
|
|
# Space limit is set as 10 GiB
|
|
with patch(
|
|
"zerver.models.Realm.currently_used_upload_space_bytes",
|
|
return_value=11 * 1024 * 1024 * 1024,
|
|
):
|
|
result = self.client_post("/json/export/realm")
|
|
self.assert_json_error(
|
|
result,
|
|
f"Please request a manual export from {settings.ZULIP_ADMINISTRATOR}.",
|
|
)
|
|
|
|
# Message limit is set as 250000
|
|
realm_count.value = 250001
|
|
realm_count.save(update_fields=["value"])
|
|
result = self.client_post("/json/export/realm")
|
|
self.assert_json_error(
|
|
result,
|
|
f"Please request a manual export from {settings.ZULIP_ADMINISTRATOR}.",
|
|
)
|