def test_fails_if_job_failed(self):
        solid_config = {
            "solids": {
                "base_wait_for_job_completion": {
                    "config": {
                        "poll_interval_seconds": 0,
                        "max_wait_time_seconds": 10,
                    }
                }
            }
        }

        with patch(
                'dagster_utils.resources.data_repo.jade_data_repo.NoopDataRepoClient.retrieve_job',
                return_value=mock_job_status(
                    completed=True, successful=False)) as mocked_retrieve_job:
            with self.assertRaisesRegex(Failure,
                                        "Job did not complete successfully."):
                result = execute_solid(base_wait_for_job_completion,
                                       run_config=solid_config,
                                       mode_def=self.test_mode,
                                       input_values={
                                           'job_id': JobId('steve-was-here'),
                                       })
                self.assertTrue(result.success)
                self.assertEqual(mocked_retrieve_job.call_count, 3)
    def test_check_data_ingest_job_retries_on_5xx(self):
        data_repo = Mock(spec=RepositoryApi)
        api_responses = [ApiException(status=502), {'failedFiles': 0}]
        data_repo.retrieve_job_result = Mock(side_effect=api_responses)
        mode_def = ModeDefinition(
            name='test',
            resource_defs={
                "data_repo_client":
                ResourceDefinition.hardcoded_resource(data_repo)
            })

        result: SolidExecutionResult = execute_solid(
            base_check_data_ingest_job_result,
            mode_def=mode_def,
            input_values={'job_id': JobId('fake_job_id')},
            run_config={
                'solids': {
                    'base_check_data_ingest_job_result': {
                        'config': {
                            'max_wait_time_seconds': 3,
                            'poll_interval_seconds': 1
                        }
                    }
                }
            })

        self.assertTrue(result.success,
                        "Poll ingest should not raise after a single 5xx")
    def test_polls_repeatedly_until_complete(self):
        solid_config = {
            "solids": {
                "base_wait_for_job_completion": {
                    "config": {
                        "poll_interval_seconds": 0,
                        "max_wait_time_seconds": 10,
                    }
                }
            }
        }
        job_status_sequence = [
            mock_job_status(False),
            mock_job_status(False),
            mock_job_status(True)
        ]

        with patch(
                'dagster_utils.resources.data_repo.jade_data_repo.NoopDataRepoClient.retrieve_job',
                side_effect=job_status_sequence) as mocked_retrieve_job:
            result = execute_solid(base_wait_for_job_completion,
                                   run_config=solid_config,
                                   mode_def=self.test_mode,
                                   input_values={
                                       'job_id': JobId('steve-was-here'),
                                   })
            self.assertTrue(result.success)
            self.assertEqual(mocked_retrieve_job.call_count, 3)
 def test_submit_snapshot_job_calls_submit_snapshot_job_in_hca_manage(self):
     with patch(
             'hca_manage.snapshot.SnapshotManager.submit_snapshot_request_with_name',
             return_value=JobId("abcde")) as submit_snap:
         result = execute_solid(
             submit_snapshot_job,
             mode_def=test_mode,
             run_config={
                 'resources': {
                     'snapshot_config': {
                         'config': {
                             'dataset_name': 'badset',
                             'snapshot_name': 'namityname'
                         }
                     }
                 }
             },
         )
         self.assertTrue(result.success)
         self.assertEqual(result.output_value(), JobId('abcde'))
         submit_snap.assert_called_once_with('namityname')
    def test_check_data_ingest_job_result_failed_files(self):
        data_repo = Mock(spec=RepositoryApi)
        job_response = {"failedFiles": 1}
        data_repo.retrieve_job_result = Mock(return_value=job_response)
        mode_def = ModeDefinition(
            name='test',
            resource_defs={
                "data_repo_client":
                ResourceDefinition.hardcoded_resource(data_repo)
            })

        with self.assertRaises(
                Failure,
                msg="Bulk ingest should fail if a file fails to ingest"):
            execute_solid(check_data_ingest_job_result,
                          mode_def=mode_def,
                          input_values={'job_id': JobId('fake_job_id')})
    def test_check_data_ingest_job_result_should_poll_jade(self):
        data_repo = Mock(spec=RepositoryApi)
        job_response = {"failedFiles": 0}
        data_repo.retrieve_job_result = Mock(return_value=job_response)
        mode_def = ModeDefinition(
            name='test',
            resource_defs={
                "data_repo_client":
                ResourceDefinition.hardcoded_resource(data_repo)
            })

        result: SolidExecutionResult = execute_solid(
            check_data_ingest_job_result,
            mode_def=mode_def,
            input_values={'job_id': JobId('fake_job_id')})

        self.assertTrue(result.success)
    def test_fails_if_max_time_exceeded(self):
        solid_config = {
            "solids": {
                "base_wait_for_job_completion": {
                    "config": {
                        "poll_interval_seconds": 1,
                        "max_wait_time_seconds": 1,
                    }
                }
            }
        }

        with patch(
                'dagster_utils.resources.data_repo.jade_data_repo.NoopDataRepoClient.retrieve_job',
                return_value=mock_job_status(False)):
            with self.assertRaisesRegex(Failure, "Exceeded max wait time"):
                execute_solid(base_wait_for_job_completion,
                              run_config=solid_config,
                              mode_def=self.test_mode,
                              input_values={
                                  'job_id': JobId('steve-was-here'),
                              })