Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 60 additions & 12 deletions tools-tests/upload_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,27 @@ def main_norun_redacted_zip(main_norun):
yield


@pytest.fixture
def taskrunner_workdir(tmp_path: pathlib.Path):
"""
Creates a temporary workdir for the TaskRunner to use. The directory is required to exist before TaskRunner is
instantiated and closed when TaskRunner goes out of scope.
"""
workdir = tmp_path / "workdir"
workdir.mkdir()
yield workdir


class FakeResponse:
def __init__(self, status_code):
self.status_code = status_code

def getcode(self):
return self.status_code

def read(self) -> bytes:
return b"{}"


class FakeSuccessUrlOpener:
def __init__(self, *args, **kwargs):
Expand All @@ -66,28 +80,40 @@ def open(self, *args, **kwargs):

@pytest.mark.usefixtures("main_norun")
@pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]])
def test_main_output_exists(args):
def test_main_output_exists(args, taskrunner_workdir):
with pytest.raises(SystemExit, check=lambda e: e.code == 0):
with unittest.mock.patch("sys.argv", ["sg_collect", *args, ZIP_NAME]):
with unittest.mock.patch(
"sys.argv", ["sg_collect", *args, "--tmp-dir", taskrunner_workdir, ZIP_NAME]
):
sgcollect.main()
assert pathlib.Path(ZIP_NAME).exists()
assert not pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun_redacted_zip")
def test_main_output_exists_with_redacted():
def test_main_output_exists_with_redacted(taskrunner_workdir):
with pytest.raises(SystemExit, check=lambda e: e.code == 0):
with unittest.mock.patch(
"sys.argv", ["sg_collect", "--log-redaction-level", "partial", ZIP_NAME]
"sys.argv",
[
"sg_collect",
"--log-redaction-level",
"partial",
"--tmp-dir",
taskrunner_workdir,
ZIP_NAME,
],
):
sgcollect.main()
assert pathlib.Path(ZIP_NAME).exists()
assert pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun")
@pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]])
def test_main_zip_deleted_on_upload_success(args):
def test_main_zip_deleted_on_upload_success(args, taskrunner_workdir):
with unittest.mock.patch("tasks.urllib.request.build_opener", FakeSuccessUrlOpener):
with unittest.mock.patch(
"sys.argv",
Expand All @@ -98,6 +124,8 @@ def test_main_zip_deleted_on_upload_success(args):
"https://example.com",
"--customer",
"fakeCustomer",
"--tmp-dir",
taskrunner_workdir,
ZIP_NAME,
],
):
Expand All @@ -106,11 +134,12 @@ def test_main_zip_deleted_on_upload_success(args):
assert exc.value.code == 0
assert not pathlib.Path(ZIP_NAME).exists()
assert not pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun")
@pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]])
def test_main_zip_deleted_on_upload_failure(args):
def test_main_zip_deleted_on_upload_failure(args, taskrunner_workdir):
with unittest.mock.patch("tasks.urllib.request.build_opener", FakeFailureUrlOpener):
with unittest.mock.patch(
"sys.argv",
Expand All @@ -129,10 +158,11 @@ def test_main_zip_deleted_on_upload_failure(args):
assert exc.value.code == 1
assert not pathlib.Path(ZIP_NAME).exists()
assert not pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun_redacted_zip")
def test_main_redacted_zip_deleted_on_upload_success():
def test_main_redacted_zip_deleted_on_upload_success(taskrunner_workdir):
with unittest.mock.patch("tasks.urllib.request.build_opener", FakeSuccessUrlOpener):
with unittest.mock.patch(
"sys.argv",
Expand All @@ -144,6 +174,8 @@ def test_main_redacted_zip_deleted_on_upload_success():
"https://example.com",
"--customer",
"fakeCustomer",
"--tmp-dir",
taskrunner_workdir,
ZIP_NAME,
],
):
Expand All @@ -152,10 +184,11 @@ def test_main_redacted_zip_deleted_on_upload_success():
assert exc.value.code == 0
assert not pathlib.Path(ZIP_NAME).exists()
assert not pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun_redacted_zip")
def test_main_redacted_zip_deleted_on_upload_failure():
def test_main_redacted_zip_deleted_on_upload_failure(taskrunner_workdir):
with unittest.mock.patch("tasks.urllib.request.build_opener", FakeFailureUrlOpener):
with unittest.mock.patch(
"sys.argv",
Expand All @@ -167,6 +200,8 @@ def test_main_redacted_zip_deleted_on_upload_failure():
"https://example.com",
"--customer",
"fakeCustomer",
"--tmp-dir",
taskrunner_workdir,
ZIP_NAME,
],
):
Expand All @@ -175,11 +210,12 @@ def test_main_redacted_zip_deleted_on_upload_failure():
assert exc.value.code == 1
assert not pathlib.Path(ZIP_NAME).exists()
assert not pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun")
@pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]])
def test_main_keep_zip_on_upload_success(args):
def test_main_keep_zip_on_upload_success(args, taskrunner_workdir):
with unittest.mock.patch("tasks.urllib.request.build_opener", FakeSuccessUrlOpener):
with unittest.mock.patch(
"sys.argv",
Expand All @@ -191,6 +227,8 @@ def test_main_keep_zip_on_upload_success(args):
"--customer",
"fakeCustomer",
"--keep-zip",
"--tmp-dir",
taskrunner_workdir,
ZIP_NAME,
],
):
Expand All @@ -199,11 +237,12 @@ def test_main_keep_zip_on_upload_success(args):
assert exc.value.code == 0
assert pathlib.Path(ZIP_NAME).exists()
assert not pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun")
@pytest.mark.parametrize("args", [[], ["--log-redaction-level", "none"]])
def test_main_keep_zip_on_upload_failure(args):
def test_main_keep_zip_on_upload_failure(args, taskrunner_workdir):
with unittest.mock.patch("tasks.urllib.request.build_opener", FakeFailureUrlOpener):
with unittest.mock.patch(
"sys.argv",
Expand All @@ -215,6 +254,8 @@ def test_main_keep_zip_on_upload_failure(args):
"--customer",
"fakeCustomer",
"--keep-zip",
"--tmp-dir",
taskrunner_workdir,
ZIP_NAME,
],
):
Expand All @@ -223,10 +264,11 @@ def test_main_keep_zip_on_upload_failure(args):
assert exc.value.code == 1
assert pathlib.Path(ZIP_NAME).exists()
assert not pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun_redacted_zip")
def test_main_keep_zip_deleted_on_upload_success():
def test_main_keep_zip_deleted_on_upload_success(taskrunner_workdir):
with unittest.mock.patch("tasks.urllib.request.build_opener", FakeSuccessUrlOpener):
with unittest.mock.patch(
"sys.argv",
Expand All @@ -239,6 +281,8 @@ def test_main_keep_zip_deleted_on_upload_success():
"--customer",
"fakeCustomer",
"--keep-zip",
"--tmp-dir",
taskrunner_workdir,
ZIP_NAME,
],
):
Expand All @@ -247,10 +291,11 @@ def test_main_keep_zip_deleted_on_upload_success():
assert exc.value.code == 0
assert pathlib.Path(ZIP_NAME).exists()
assert pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.mark.usefixtures("main_norun_redacted_zip")
def test_main_keep_zip_deleted_on_upload_failure():
def test_main_keep_zip_deleted_on_upload_failure(taskrunner_workdir):
with unittest.mock.patch("tasks.urllib.request.build_opener", FakeFailureUrlOpener):
with unittest.mock.patch(
"sys.argv",
Expand All @@ -263,6 +308,8 @@ def test_main_keep_zip_deleted_on_upload_failure():
"--customer",
"fakeCustomer",
"--keep-zip",
"--tmp-dir",
taskrunner_workdir,
ZIP_NAME,
],
):
Expand All @@ -271,6 +318,7 @@ def test_main_keep_zip_deleted_on_upload_failure():
assert exc.value.code == 1
assert pathlib.Path(ZIP_NAME).exists()
assert pathlib.Path(REDACTED_ZIP_NAME).exists()
assert not [x for x in taskrunner_workdir.iterdir()]


@pytest.fixture(scope="session")
Expand Down
118 changes: 58 additions & 60 deletions tools/sgcollect.py
Original file line number Diff line number Diff line change
Expand Up @@ -927,77 +927,75 @@ def main() -> NoReturn:
do_upload(args[0], options.just_upload_into, options.upload_proxy)

# Create a TaskRunner and run all of the OS tasks (collect top, netstat, etc)
# The output of the tasks will go directly into couchbase.log
runner = TaskRunner(
with TaskRunner(
verbosity=options.verbosity,
default_name="sync_gateway.log",
tmp_dir=options.tmp_dir,
)

if not options.product_only:
for task in make_os_tasks(["sync_gateway"]):
runner.run(task)

# Output the Python version if verbosity was enabled
if options.verbosity:
log("Python version: %s" % sys.version)

# Find path to sg binary
sg_binary_path = discover_sg_binary_path(options, sg_url, auth_headers)

# Run SG specific tasks
for task in make_sg_tasks(
sg_url=sg_url,
auth_headers=auth_headers,
sync_gateway_config_path_option=options.sync_gateway_config,
sync_gateway_executable_path=options.sync_gateway_executable,
should_redact=should_redact,
):
runner.run(task)

if (
sg_binary_path is not None
and sg_binary_path != ""
and os.path.exists(sg_binary_path)
):
runner.collect_file(sg_binary_path)
else:
print(
"WARNING: unable to find Sync Gateway executable, omitting from result. Go pprofs will not be accurate."
)
) as runner:
if not options.product_only:
for task in make_os_tasks(["sync_gateway"]):
runner.run(task)

runner.run(get_sgcollect_info_options_task(options, args))
# Output the Python version if verbosity was enabled
if options.verbosity:
log("Python version: %s" % sys.version)

runner.close_all_files()
# Find path to sg binary
sg_binary_path = discover_sg_binary_path(options, sg_url, auth_headers)

# Build redacted zip file
if should_redact:
log("Redacting log files to level: %s" % options.redact_level)
runner.redact_and_zip(
redact_zip_file, "sgcollect_info", options.salt_value, platform.node()
)
# Run SG specific tasks
for task in make_sg_tasks(
sg_url=sg_url,
auth_headers=auth_headers,
sync_gateway_config_path_option=options.sync_gateway_config,
sync_gateway_executable_path=options.sync_gateway_executable,
should_redact=should_redact,
):
runner.run(task)

# Build the actual zip file
runner.zip(zip_filename, "sgcollect_info", platform.node())
if (
sg_binary_path is not None
and sg_binary_path != ""
and os.path.exists(sg_binary_path)
):
runner.collect_file(sg_binary_path)
else:
print(
"WARNING: unable to find Sync Gateway executable, omitting from result. Go pprofs will not be accurate."
)

if options.redact_level != "none":
print("Zipfile built: {0}".format(redact_zip_file))
runner.run(get_sgcollect_info_options_task(options, args))

print("Zipfile built: {0}".format(zip_filename))
runner.close_all_files()

if not upload_url:
sys.exit(0)
# Upload the zip to the URL to S3 if required
try:
# Build redacted zip file
if should_redact:
exit_code = do_upload(redact_zip_file, upload_url, options.upload_proxy)
else:
exit_code = do_upload(zip_filename, upload_url, options.upload_proxy)
finally:
if not options.keep_zip:
delete_zip(zip_filename)
delete_zip(redact_zip_file)
sys.exit(exit_code)
log("Redacting log files to level: %s" % options.redact_level)
runner.redact_and_zip(
redact_zip_file, "sgcollect_info", options.salt_value, platform.node()
)

try:
# Build the actual zip file
runner.zip(zip_filename, "sgcollect_info", platform.node())

if options.redact_level != "none":
print("Zipfile built: {0}".format(redact_zip_file))

print("Zipfile built: {0}".format(zip_filename))

if not upload_url:
sys.exit(0)
# Upload the zip to the URL to S3 if required
if should_redact:
exit_code = do_upload(redact_zip_file, upload_url, options.upload_proxy)
else:
exit_code = do_upload(zip_filename, upload_url, options.upload_proxy)
sys.exit(exit_code)
finally:
if not options.keep_zip and upload_url:
Copy link

Copilot AI Jan 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The condition not options.keep_zip and upload_url will prevent zip deletion when upload_url is empty/None, even without the --keep-zip flag. This changes the behavior: previously, zips were deleted in the finally block regardless of upload_url when --keep-zip was not set. Now they're only deleted when an upload URL exists. This could leave zip files behind when no upload is performed. Consider if this behavior change is intended or if the condition should be if not options.keep_zip:.

Suggested change
if not options.keep_zip and upload_url:
if not options.keep_zip:

Copilot uses AI. Check for mistakes.
delete_zip(zip_filename)
delete_zip(redact_zip_file)


def ud(value, should_redact=True):
Expand Down
Loading
Loading