diff --git a/dnastack/cli/commands/workbench/runs/samples/commands.py b/dnastack/cli/commands/workbench/runs/samples/commands.py index 6727c328..47ec9a1a 100644 --- a/dnastack/cli/commands/workbench/runs/samples/commands.py +++ b/dnastack/cli/commands/workbench/runs/samples/commands.py @@ -50,8 +50,18 @@ def add_samples(context: Optional[str] = None, run = client.get_run(run_id) existing = run.request.samples if run.request and run.request.samples else [] + new_samples = [] + provided_sample_ids = set(samples) + + # Update existing samples with new storage_account_id if provided + for s in existing: + if s.id in provided_sample_ids and storage_account_id is not None: + new_samples.append(SimpleSample(id=s.id, storage_account_id=storage_account_id)) + else: + new_samples.append(s) + + # Add truly new samples existing_ids = {s.id for s in existing} - new_samples = list(existing) for sample_id in samples: if sample_id not in existing_ids: new_samples.append(SimpleSample(id=sample_id, storage_account_id=storage_account_id)) diff --git a/tests/unit/cli/workbench/test_runs_samples_commands.py b/tests/unit/cli/workbench/test_runs_samples_commands.py index bc9e727e..9bc7513a 100644 --- a/tests/unit/cli/workbench/test_runs_samples_commands.py +++ b/tests/unit/cli/workbench/test_runs_samples_commands.py @@ -106,6 +106,65 @@ def test_add_samples_missing_run_id(self, mock_get_client): self.assertNotEqual(result.exit_code, 0) + @patch('dnastack.cli.commands.workbench.runs.samples.commands.get_ewes_client') + def test_add_existing_sample_updates_storage_account(self, mock_get_client): + """Test that re-adding existing samples with --storage-account updates their storage_account_id""" + mock_get_client.return_value = self.mock_ewes_client + + existing_run = Mock(spec=ExtendedRun) + existing_run.request = Mock(spec=ExtendedRunRequest) + existing_run.request.samples = [ + SimpleSample(id='sample-1', storage_account_id=None), + SimpleSample(id='sample-2', storage_account_id=None), + ] + self.mock_ewes_client.get_run.return_value = existing_run + + mock_result = Mock(spec=ExtendedRunStatus) + mock_result.model_dump.return_value = {'samples': []} + self.mock_ewes_client.update_run_samples.return_value = mock_result + + result = self.runner.invoke( + self.group, + ['add', '--run-id', 'run-1', '--sample', 'sample-1', '--sample', 'sample-2', '--storage-account', 'sa-new'] + ) + + self.assertEqual(result.exit_code, 0) + call_args = self.mock_ewes_client.update_run_samples.call_args + samples = call_args[0][1] + self.assertEqual(len(samples), 2) + for s in samples: + self.assertEqual(s.storage_account_id, 'sa-new', + f"Sample {s.id} should have updated storage_account_id") + + @patch('dnastack.cli.commands.workbench.runs.samples.commands.get_ewes_client') + def test_add_existing_sample_preserves_storage_account_when_not_provided(self, mock_get_client): + """Test that re-adding existing samples WITHOUT --storage-account preserves their existing storage_account_id""" + mock_get_client.return_value = self.mock_ewes_client + + existing_run = Mock(spec=ExtendedRun) + existing_run.request = Mock(spec=ExtendedRunRequest) + existing_run.request.samples = [ + SimpleSample(id='sample-1', storage_account_id='sa-original'), + ] + self.mock_ewes_client.get_run.return_value = existing_run + + mock_result = Mock(spec=ExtendedRunStatus) + mock_result.model_dump.return_value = {'samples': []} + self.mock_ewes_client.update_run_samples.return_value = mock_result + + result = self.runner.invoke( + self.group, + ['add', '--run-id', 'run-1', '--sample', 'sample-1'] + ) + + self.assertEqual(result.exit_code, 0) + call_args = self.mock_ewes_client.update_run_samples.call_args + samples = call_args[0][1] + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0].storage_account_id, 'sa-original', + "Existing storage_account_id should be preserved when --storage-account not provided") + + class TestRunsSamplesRemoveCommand(unittest.TestCase): """Unit tests for runs samples remove command"""