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()
