inclusion of more data row get tests

Author: jtsodapop

tests/integration/test_global_keys.py

@@ -145,3 +145,111 @@ def test_global_key_with_whitespaces_validation(client, dataset, image_url):
     assert assign_errors_gks == set([gk_1, gk_2, gk_3])
     assert assign_errors_msgs == set(
         ['Invalid global key', 'Invalid global key', 'Invalid global key'])
+
+
+def test_get_data_rows_for_global_keys(client, dataset, image_url):
+    gk_1 = str(uuid.uuid4())
+    gk_2 = str(uuid.uuid4())
+
+    dr_1 = dataset.create_data_row(row_data=image_url,
+                                   external_id="hello",
+                                   global_key=gk_1)
+    dr_2 = dataset.create_data_row(row_data=image_url,
+                                   external_id="world",
+                                   global_key=gk_2)
+
+    res = client.get_data_rows_for_global_keys([gk_1])
+    assert res['status'] == "SUCCESS"
+    assert res['errors'] == []
+    assert res['results'] == [dr_1]
+
+    res = client.get_data_rows_for_global_keys([gk_2])
+    assert res['status'] == "SUCCESS"
+    assert res['errors'] == []
+    assert res['results'] == [dr_2]
+
+    res = client.get_data_rows_for_global_keys([gk_1, gk_2])
+    assert res['status'] == "SUCCESS"
+    assert res['errors'] == []
+    assert res['results'] == [dr_1, dr_2]
+
+
+def test_get_data_rows_for_invalid_global_keys(client, dataset, image_url):
+    gk_1 = str(uuid.uuid4())
+    gk_2 = str(uuid.uuid4())
+
+    dr_1 = dataset.create_data_row(row_data=image_url, external_id="hello")
+    dr_2 = dataset.create_data_row(row_data=image_url,
+                                   external_id="world",
+                                   global_key=gk_2)
+
+    res = client.get_data_rows_for_global_keys([gk_1])
+    assert res['status'] == "FAILURE"
+    assert len(res['errors']) == 1
+    assert res['errors'][0]['error'] == "Data Row not found"
+    assert res['errors'][0]['global_key'] == gk_1
+
+    res = client.get_data_rows_for_global_keys([gk_1, gk_2])
+    assert res['status'] == "PARTIAL SUCCESS"
+
+    assert len(res['errors']) == 1
+    assert len(res['results']) == 1
+
+    assert res['errors'][0]['error'] == "Data Row not found"
+    assert res['errors'][0]['global_key'] == gk_1
+
+    assert res['results'][0] == dr_2
+
+
+def test_get_data_row_ids_for_global_keys(client, dataset, image_url):
+    gk_1 = str(uuid.uuid4())
+    gk_2 = str(uuid.uuid4())
+
+    dr_1 = dataset.create_data_row(row_data=image_url,
+                                   external_id="hello",
+                                   global_key=gk_1)
+    dr_2 = dataset.create_data_row(row_data=image_url,
+                                   external_id="world",
+                                   global_key=gk_2)
+
+    res = client.get_data_row_ids_for_global_keys([gk_1])
+    assert res['status'] == "SUCCESS"
+    assert res['errors'] == []
+    assert res['results'] == [dr_1.uid]
+
+    res = client.get_data_row_ids_for_global_keys([gk_2])
+    assert res['status'] == "SUCCESS"
+    assert res['errors'] == []
+    assert res['results'] == [dr_2.uid]
+
+    res = client.get_data_row_ids_for_global_keys([gk_1, gk_2])
+    assert res['status'] == "SUCCESS"
+    assert res['errors'] == []
+    assert res['results'] == [dr_1.uid, dr_2.uid]
+
+
+def test_get_data_row_ids_for_invalid_global_keys(client, dataset, image_url):
+    gk_1 = str(uuid.uuid4())
+    gk_2 = str(uuid.uuid4())
+
+    dr_1 = dataset.create_data_row(row_data=image_url, external_id="hello")
+    dr_2 = dataset.create_data_row(row_data=image_url,
+                                   external_id="world",
+                                   global_key=gk_2)
+
+    res = client.get_data_row_ids_for_global_keys([gk_1])
+    assert res['status'] == "FAILURE"
+    assert len(res['errors']) == 1
+    assert res['errors'][0]['error'] == "Data Row not found"
+    assert res['errors'][0]['global_key'] == gk_1
+
+    res = client.get_data_row_ids_for_global_keys([gk_1, gk_2])
+    assert res['status'] == "PARTIAL SUCCESS"
+
+    assert len(res['errors']) == 1
+    assert len(res['results']) == 1
+
+    assert res['errors'][0]['error'] == "Data Row not found"
+    assert res['errors'][0]['global_key'] == gk_1
+
+    assert res['results'][0] == dr_2.uid