def setUp(self) -> None: self.job_type = "test" self.builder = DataProcJobBuilder( project_id=GCP_PROJECT, task_id=TASK_ID, cluster_name=CLUSTER_NAME, job_type=self.job_type, properties={"test": "test"}, )
def test_init(self, mock_uuid): mock_uuid.return_value = "uuid" properties = {"test": "test"} job = { "job": { "labels": { "airflow-version": AIRFLOW_VERSION }, "placement": { "cluster_name": CLUSTER_NAME }, "reference": { "job_id": TASK_ID + "_uuid", "project_id": GCP_PROJECT }, "test": { "properties": properties }, } } builder = DataProcJobBuilder( project_id=GCP_PROJECT, task_id=TASK_ID, cluster_name=CLUSTER_NAME, job_type="test", properties=properties, ) assert job == builder.job
class TestDataProcJobBuilder(unittest.TestCase): def setUp(self) -> None: self.job_type = "test" self.builder = DataProcJobBuilder( project_id=GCP_PROJECT, task_id=TASK_ID, cluster_name=CLUSTER_NAME, job_type=self.job_type, properties={"test": "test"}, ) @mock.patch(DATAPROC_STRING.format("uuid.uuid4")) def test_init(self, mock_uuid): mock_uuid.return_value = "uuid" properties = {"test": "test"} job = { "job": { "labels": { "airflow-version": AIRFLOW_VERSION }, "placement": { "cluster_name": CLUSTER_NAME }, "reference": { "job_id": TASK_ID + "_uuid", "project_id": GCP_PROJECT }, "test": { "properties": properties }, } } builder = DataProcJobBuilder( project_id=GCP_PROJECT, task_id=TASK_ID, cluster_name=CLUSTER_NAME, job_type="test", properties=properties, ) assert job == builder.job def test_add_labels(self): labels = {"key": "value"} self.builder.add_labels(labels) assert "key" in self.builder.job["job"]["labels"] assert "value" == self.builder.job["job"]["labels"]["key"] def test_add_variables(self): variables = ["variable"] self.builder.add_variables(variables) assert variables == self.builder.job["job"][ self.job_type]["script_variables"] def test_add_args(self): args = ["args"] self.builder.add_args(args) assert args == self.builder.job["job"][self.job_type]["args"] def test_add_query(self): query = ["query"] self.builder.add_query(query) assert { "queries": [query] } == self.builder.job["job"][self.job_type]["query_list"] def test_add_query_uri(self): query_uri = "query_uri" self.builder.add_query_uri(query_uri) assert query_uri == self.builder.job["job"][ self.job_type]["query_file_uri"] def test_add_jar_file_uris(self): jar_file_uris = ["jar_file_uris"] self.builder.add_jar_file_uris(jar_file_uris) assert jar_file_uris == self.builder.job["job"][ self.job_type]["jar_file_uris"] def test_add_archive_uris(self): archive_uris = ["archive_uris"] self.builder.add_archive_uris(archive_uris) assert archive_uris == self.builder.job["job"][ self.job_type]["archive_uris"] def test_add_file_uris(self): file_uris = ["file_uris"] self.builder.add_file_uris(file_uris) assert file_uris == self.builder.job["job"][self.job_type]["file_uris"] def test_add_python_file_uris(self): python_file_uris = ["python_file_uris"] self.builder.add_python_file_uris(python_file_uris) assert python_file_uris == self.builder.job["job"][ self.job_type]["python_file_uris"] def test_set_main_error(self): with pytest.raises(Exception): self.builder.set_main("test", "test") def test_set_main_class(self): main = "main" self.builder.set_main(main_class=main, main_jar=None) assert main == self.builder.job["job"][self.job_type]["main_class"] def test_set_main_jar(self): main = "main" self.builder.set_main(main_class=None, main_jar=main) assert main == self.builder.job["job"][ self.job_type]["main_jar_file_uri"] def test_set_python_main(self): main = "main" self.builder.set_python_main(main) assert main == self.builder.job["job"][ self.job_type]["main_python_file_uri"] @mock.patch(DATAPROC_STRING.format("uuid.uuid4")) def test_set_job_name(self, mock_uuid): uuid = "test_uuid" mock_uuid.return_value = uuid name = "name" self.builder.set_job_name(name) name += "_" + uuid[:8] assert name == self.builder.job["job"]["reference"]["job_id"] def test_build(self): assert self.builder.job == self.builder.build()