mirror of
https://github.com/zulip/zulip.git
synced 2025-11-06 15:03:34 +00:00
76bcb96414
As suggested by the new comments, the cost for a Zulip data export scales with messages actually included in the export, so an organizations with 1M private messages but only 50K public stream messages should not be modeled the same as one with 1M public stream messages for the purpose of the limits here. Also improve the comments and variable names more generally.
232 lines
9.2 KiB
Python
232 lines
9.2 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 (
|
|
HostRequestMock,
|
|
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(HostRequestMock(), 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()
|
|
assert audit_log_entry is not None
|
|
self.assertEqual(audit_log_entry.acting_user_id, admin.id)
|
|
|
|
# Test that the file is hosted, and the contents are as expected.
|
|
extra_data = audit_log_entry.extra_data
|
|
assert extra_data is not None
|
|
export_path = orjson.loads(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")
|
|
response_dict = self.assert_json_success(result)
|
|
|
|
# Test that the export we have is the export we created.
|
|
export_dict = response_dict["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()
|
|
extra_data = audit_log_entry.extra_data
|
|
assert extra_data is not None
|
|
export_data = orjson.loads(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()
|
|
assert audit_log_entry is not None
|
|
self.assertEqual(audit_log_entry.acting_user_id, admin.id)
|
|
|
|
# Test that the file is hosted, and the contents are as expected.
|
|
extra_data = audit_log_entry.extra_data
|
|
assert extra_data is not None
|
|
export_path = orjson.loads(extra_data).get("export_path")
|
|
response = self.client_get(export_path)
|
|
self.assertEqual(response.status_code, 200)
|
|
self.assert_streaming_content(response, b"zulip!")
|
|
|
|
result = self.client_get("/json/export/realm")
|
|
response_dict = self.assert_json_success(result)
|
|
|
|
# Test that the export we have is the export we created.
|
|
export_dict = response_dict["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()
|
|
extra_data = audit_log_entry.extra_data
|
|
assert extra_data is not None
|
|
export_data = orjson.loads(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(HostRequestMock(), 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(),
|
|
value=0,
|
|
property="messages_sent:message_type:day",
|
|
subgroup="public_stream",
|
|
)
|
|
|
|
# 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}.",
|
|
)
|