diff --git a/launchable/commands/record/attachment.py b/launchable/commands/record/attachment.py index b1f05bf7d..b62fc89da 100644 --- a/launchable/commands/record/attachment.py +++ b/launchable/commands/record/attachment.py @@ -1,7 +1,8 @@ +import fnmatch import tarfile import zipfile from io import BytesIO -from typing import Optional +from typing import Optional, Tuple import click from tabulate import tabulate @@ -23,12 +24,20 @@ class AttachmentStatus: help='In the format builds//test_sessions/', type=str, ) +@click.option( + '--include', + 'include_patterns', + help='Include only files matching pattern (e.g., "*.log"). Can be specified multiple times.', + type=str, + multiple=True, +) @click.argument('attachments', nargs=-1) # type=click.Path(exists=True) @click.pass_context def attachment( context: click.core.Context, attachments, - session: Optional[str] = None + session: Optional[str] = None, + include_patterns: Tuple[str, ...] = () ): client = LaunchableClient(app=context.obj) summary_rows = [] @@ -44,6 +53,9 @@ def attachment( if zip_info.is_dir(): continue + if not matches_include_patterns(zip_info.filename, include_patterns): + continue + file_content = zip_file.read(zip_info.filename) if not valid_utf8_file(file_content): @@ -62,6 +74,9 @@ def attachment( if tar_info.isdir(): continue + if not matches_include_patterns(tar_info.name, include_patterns): + continue + file_obj = tar_file.extractfile(tar_info) if file_obj is None: continue @@ -95,6 +110,21 @@ def attachment( display_summary_as_table(summary_rows) +def matches_include_patterns(filename: str, include_patterns: Tuple[str, ...]) -> bool: + """ + Check if a file should be included based on the include patterns. + If no patterns are specified, all files are included. + """ + if not include_patterns: + return True + + for pattern in include_patterns: + if fnmatch.fnmatch(filename, pattern): + return True + + return False + + def valid_utf8_file(file_content: bytes) -> bool: # Check for null bytes (binary files) if b'\x00' in file_content: diff --git a/tests/commands/record/test_attachment.py b/tests/commands/record/test_attachment.py index 4a4743c5d..38f8fe609 100644 --- a/tests/commands/record/test_attachment.py +++ b/tests/commands/record/test_attachment.py @@ -111,3 +111,53 @@ def test_attachment_zip_file(self): result = self.cli("record", "attachment", "--session", self.session, tar_path) self.assertIn(expect, result.output) + + @responses.activate + @mock.patch.dict(os.environ, {"LAUNCHABLE_TOKEN": CliTestCase.launchable_token}) + def test_attachment_with_include_filter(self): + with tempfile.TemporaryDirectory() as temp_dir: + # Create temporary files + text_file_1 = os.path.join(temp_dir, "app.log") + text_file_2 = os.path.join(temp_dir, "nested", "debug.log") + text_file_3 = os.path.join(temp_dir, "config.yml") + zip_path = os.path.join(temp_dir, "logs.zip") + + # Create directory structure + os.makedirs(os.path.dirname(text_file_2)) + + # Write test content + with open(text_file_1, 'w') as f: + f.write("[INFO] Test log entry") + with open(text_file_2, 'w') as f: + f.write("[DEBUG] Nested log entry") + with open(text_file_3, 'w') as f: + f.write("config: value") + + # Create zip file + with zipfile.ZipFile(zip_path, 'w') as zf: + zf.write(text_file_1, 'app.log') + zf.write(text_file_2, 'nested/debug.log') + zf.write(text_file_3, 'config.yml') + + responses.add( + responses.POST, + "{}/intake/organizations/{}/workspaces/{}/builds/{}/test_sessions/{}/attachment".format( + get_base_url(), self.organization, self.workspace, self.build_name, self.session_id), + match=[responses.matchers.header_matcher({"Content-Disposition": 'attachment;filename="app.log"'})], + status=200) + + responses.add( + responses.POST, + "{}/intake/organizations/{}/workspaces/{}/builds/{}/test_sessions/{}/attachment".format( + get_base_url(), self.organization, self.workspace, self.build_name, self.session_id), + match=[responses.matchers.header_matcher({"Content-Disposition": 'attachment;filename="nested/debug.log"'})], + status=200) + + result = self.cli("record", "attachment", "--session", self.session, "--include", "*.log", zip_path) + + expect = """| File | Status | +|------------------|-------------------------| +| app.log | ✓ Recorded successfully | +| nested/debug.log | ✓ Recorded successfully | +""" + self.assertIn(expect, result.output)