Repository: spotify/luigi
File: test/contrib/bigquery_gcloud_test.py (Python)

Code before:
                        {"name": "inner", "type": "int", "doc": "A inner field"}
                    ]
                }, "doc": "This field shall be an inner"},
                {"name": "col2", "type": "int", "doc": "The beautiful"}
            ]
        }""")
        writer = DataFileWriter(open("tmp.avro", "wb"), DatumWriter(), schema)
        writer.append({'col0': 1000, 'col1': {'inner': 1234}, 'col2': 1001})
        writer.close()
        self.gcs_client.put("tmp.avro", self.gcs_dir_url)
        # FIXME don't need a tmp.avro; use GCS API like objects().insert that accepts file contents in a string:
        # https://developers.google.com/resources/api-libraries/documentation/storage/v1/python/latest/storage_v1.objects.html#insert

    def setUp(self):
        self.gcs_client = gcs.GCSClient(CREDENTIALS)
        self.bq_client = bigquery.BigQueryClient(CREDENTIALS)

        self.table_id = "avro_bq_table"
        self.gcs_dir_url = 'gs://' + BUCKET_NAME + "/foo/tmp.avro"
        self._produce_test_input()
        self.addCleanup(self.bq_client.delete_dataset, bigquery.BQDataset(PROJECT_ID, DATASET_ID, EU_LOCATION))
        self.addCleanup(self.gcs_client.remove, self.gcs_dir_url)

    def test_load_avro_dir_and_propagate_doc(self):
        class BigQueryLoadAvroTestInput(luigi.ExternalTask):
            def output(_):
                return gcs.GCSTarget(self.gcs_dir_url)

        class BigQueryLoadAvroTestTask(bigquery_avro.BigQueryLoadAvro):
            def requires(_):
                return BigQueryLoadAvroTestInput()

            def output(_):
                return bigquery.BigQueryTarget(PROJECT_ID, DATASET_ID, self.table_id)

        task = BigQueryLoadAvroTestTask()
        self.assertFalse(task.complete())
        task.run()
        self.assertTrue(task.complete())

        table = self.bq_client.client.tables().get(projectId=PROJECT_ID,
                                                   datasetId=DATASET_ID,
                                                   tableId=self.table_id).execute()
        self.assertEqual(table['description'], 'The description')
        self.assertEqual(table['schema']['fields'][0]['description'], 'The bold')
        self.assertEqual(table['schema']['fields'][1]['description'], 'This field shall be an inner')
        self.assertEqual(table['schema']['fields'][1]['fields'][0]['description'], 'A inner field')
        self.assertEqual(table['schema']['fields'][2]['description'], 'The beautiful')


Review comment (bug):
This should come before producing test input or immediately after it, otherwise there's code in between that can crash and thus prevent adding of the cleanup.


Code after:
                        {"name": "inner", "type": "int", "doc": "A inner field"},
                        {"name": "col0", "type": "int", "doc": "Same name as outer but different doc"},
                        {"name": "col1", "type": ["null", "string"], "default": null, "doc": "Nullable primitive"},
                        {"name": "col2", "type": ["null", {
                            "type": "map",
                            "values": "string"
                        }], "default": null, "doc": "Nullable map"}
                    ]
                }, "doc": "This field shall be an inner"},
                {"name": "col2", "type": "int", "doc": "The beautiful"},
                {"name": "col3", "type": "double"}
            ]
        }""")
        self.addCleanup(os.remove, "tmp.avro")
        writer = DataFileWriter(open("tmp.avro", "wb"), DatumWriter(), schema)
        writer.append({'col0': 1000, 'col1': {'inner': 1234, 'col0': 3000}, 'col2': 1001, 'col3': 1.001})
        writer.close()
        self.gcs_client.put("tmp.avro", self.gcs_dir_url + "/tmp.avro")

    def setUp(self):
        self.gcs_client = gcs.GCSClient(CREDENTIALS)
        self.bq_client = bigquery.BigQueryClient(CREDENTIALS)

        self.table_id = "avro_bq_table"
        self.gcs_dir_url = 'gs://' + BUCKET_NAME + "/foo"
        self.addCleanup(self.gcs_client.remove, self.gcs_dir_url)
        self.addCleanup(self.bq_client.delete_dataset, bigquery.BQDataset(PROJECT_ID, DATASET_ID, EU_LOCATION))
        self._produce_test_input()

    def test_load_avro_dir_and_propagate_doc(self):
        class BigQueryLoadAvroTestInput(luigi.ExternalTask):
            def output(_):
                return gcs.GCSTarget(self.gcs_dir_url)

        class BigQueryLoadAvroTestTask(bigquery_avro.BigQueryLoadAvro):
            def requires(_):
                return BigQueryLoadAvroTestInput()

            def output(_):
                return bigquery.BigQueryTarget(PROJECT_ID, DATASET_ID, self.table_id)

        task = BigQueryLoadAvroTestTask()
        self.assertFalse(task.complete())
        task.run()
        self.assertTrue(task.complete())

        table = self.bq_client.client.tables().get(projectId=PROJECT_ID,
                                                   datasetId=DATASET_ID,
                                                   tableId=self.table_id).execute()
        self.assertEqual(table['description'], 'The description')
        self.assertEqual(table['schema']['fields'][0]['description'], 'The bold')