Merge 3dcd155 into c809f68

src/restore/test/restore_test.py

@@ -1,14 +1,16 @@
 import json
 import logging
+
 from datetime import datetime
 
 from src.commons.big_query.big_query import BigQuery
 from src.commons.big_query.big_query_table_metadata import BigQueryTableMetadata
 from src.commons.config.configuration import configuration
-from src.commons.google_cloud_storage_client import GoogleCloudStorageClient as gcs
+from src.commons.google_cloud_storage_client import \
+    GoogleCloudStorageClient as gcs
+from src.commons.table_reference import TableReference
 from src.restore.test.table_randomizer import TableRandomizer
 from src.restore.test.table_restore_invoker import TableRestoreInvoker
-from src.commons.table_reference import TableReference
 
 RESTORE_DATASET_ID_US = 'smoke_test_US'
 RESTORE_DATASET_ID_EU = 'smoke_test_EU'
@@ -32,10 +34,13 @@ def test(self, host_url):
         table_has_been_restored = response["status"]["result"] == "Success"
 
         if table_has_been_restored:
-            response = self.__check_restored_table_matches_source(response,
-                                                                  source_table)
-            BigQuery().delete_table(table_reference)
-            return response
+            result = self.__check_restored_table_matches_source(response,
+                                                                source_table)
+
+            target_table_reference = self.__get_target_tab_ref(response)
+            BigQuery().delete_table(target_table_reference)
+
+            return result
         else:
             resp_msg = "Restore test failed. " \
                        "Failed to restore a table {}".format(table_reference)
@@ -48,13 +53,9 @@ def __get_restore_dataset(src_table):
             else RESTORE_DATASET_ID_US
 
     def __check_restored_table_matches_source(self, response, src_table):
-        restoration_items = response["restorationItems"]
-        assert len(restoration_items) == 1
-
-        target_table_reference = \
-            TableReference.parse_tab_ref(restoration_items[0]['targetTable'])
-        target_table = BigQueryTableMetadata.get_table_by_reference(
-            target_table_reference)
+        target_table_reference = self.__get_target_tab_ref(response)
+        target_table = \
+            BigQueryTableMetadata.get_table_by_reference(target_table_reference)
 
         tables_match, assertion_msg = \
             self.__assert_restored_table_matches_source(src_table, target_table)
@@ -71,6 +72,12 @@ def __check_restored_table_matches_source(self, response, src_table):
             self.__save_test_status(source_table_reference, False, response)
             raise Exception(resp_msg)
 
+    @staticmethod
+    def __get_target_tab_ref(response):
+        restoration_items = response["restorationItems"]
+        assert len(restoration_items) == 1
+        return TableReference.parse_tab_ref(restoration_items[0]['targetTable'])
+
     @staticmethod
     def __pick_random_table():
         src_table = TableRandomizer().get_random_table_metadata()

tests/restore/test/test_restore_test_handler.py

@@ -1,10 +1,9 @@
 import json
 import unittest
 
+import webtest
 from apiclient.http import HttpMockSequence
 from google.appengine.ext import testbed
-
-import webtest
 from mock import patch
 
 from src.commons.big_query.big_query import BigQuery
@@ -15,7 +14,6 @@
 from src.restore.test.table_randomizer import TableRandomizer
 from src.restore.test.table_restore_invoker import TableRestoreInvoker
 
-
 example_table = {
     'tableReference': {
         'projectId': 'myproject123',
@@ -229,12 +227,6 @@ def test_handler_outputs_failed_status_if_table_was_restored_but_rows_doesnt_mat
     def test_handler_delete_restored_table_after_success(
             self, _, _1, _2, _create_http, wait_till_done, delete_table):
         # given
-        table_reference = TableReference(
-            project_id=example_table['tableReference']["projectId"],
-            dataset_id=example_table['tableReference']["datasetId"],
-            table_id=example_table['tableReference']["tableId"],
-        )
-
         endpoint_invoke_result = HttpMockSequence([(
             {'status': '200'}, json.dumps({
                 "restorationJobId": "64c6e50c-b511-43eb-ba75-f44f3d131f84"}))])
@@ -246,4 +238,6 @@ def test_handler_delete_restored_table_after_success(
 
         # then
         self.assertEquals(200, response.status_int)
-        delete_table.assert_called_with(table_reference)
+        delete_table.assert_called_with(TableReference('target-project',
+                                                       'target_dataset',
+                                                       'target_table'))