From 15fe51b788481b2a4f8d8579da229a4d0cc9b62d Mon Sep 17 00:00:00 2001
From: Alexandre Girard <alexandre@airbyte.io>
Date: Mon, 13 Jun 2022 16:51:36 -0700
Subject: [PATCH] [low-code connectors] add a few unit tests (#13666)

* add a couple of unit tests

* refactor tests
---
 .../declarative/extractors/test_jello.py      | 70 ++++++++-----------
 .../test_interpolated_boolean.py              |  3 +
 2 files changed, 32 insertions(+), 41 deletions(-)

diff --git a/airbyte-cdk/python/unit_tests/sources/declarative/extractors/test_jello.py b/airbyte-cdk/python/unit_tests/sources/declarative/extractors/test_jello.py
index c19c0fdeb725..dd1a83494190 100644
--- a/airbyte-cdk/python/unit_tests/sources/declarative/extractors/test_jello.py
+++ b/airbyte-cdk/python/unit_tests/sources/declarative/extractors/test_jello.py
@@ -4,63 +4,51 @@
 
 import json
 
+import pytest
 import requests
 from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder
 from airbyte_cdk.sources.declarative.extractors.jello import JelloExtractor
 
 config = {"field": "record_array"}
-decoder = JsonDecoder()
-
-
-def test():
-    transform = "_.data"
-    extractor = JelloExtractor(transform, decoder, config)
-
-    records = [{"id": 1}, {"id": 2}]
-    body = {"data": records}
-    response = create_response(body)
-    actual_records = extractor.extract_records(response)
-
-    assert actual_records == records
-
-
-def test_field_in_config():
-    transform = "_.{{ config['field'] }}"
-    extractor = JelloExtractor(transform, decoder, config)
+kwargs = {"data_field": "records"}
 
-    records = [{"id": 1}, {"id": 2}]
-    body = {"record_array": records}
-    response = create_response(body)
-    actual_records = extractor.extract_records(response)
-
-    assert actual_records == records
+decoder = JsonDecoder()
 
 
-def test_field_in_kwargs():
-    transform = "_.{{ kwargs['data_field'] }}"
-    kwargs = {"data_field": "records"}
+@pytest.mark.parametrize(
+    "test_name, transform, body, expected_records",
+    [
+        ("test_extract_from_array", "_.data", {"data": [{"id": 1}, {"id": 2}]}, [{"id": 1}, {"id": 2}]),
+        ("test_field_in_config", "_.{{ config['field'] }}", {"record_array": [{"id": 1}, {"id": 2}]}, [{"id": 1}, {"id": 2}]),
+        ("test_field_in_kwargs", "_.{{ kwargs['data_field'] }}", {"records": [{"id": 1}, {"id": 2}]}, [{"id": 1}, {"id": 2}]),
+        ("test_default", "_{{kwargs['field']}}", [{"id": 1}, {"id": 2}], [{"id": 1}, {"id": 2}]),
+        (
+            "test_remove_fields_from_records",
+            "[{k:v for k,v in d.items() if k != 'value_to_remove'} for d in _.data]",
+            {"data": [{"id": 1, "value": "HELLO", "value_to_remove": "fail"}, {"id": 2, "value": "WORLD", "value_to_remove": "fail"}]},
+            [{"id": 1, "value": "HELLO"}, {"id": 2, "value": "WORLD"}],
+        ),
+        (
+            "test_add_fields_from_records",
+            "[{**{k:v for k,v in d.items()}, **{'project_id': d['project']['id']}} for d in _.data]",
+            {"data": [{"id": 1, "value": "HELLO", "project": {"id": 8}}, {"id": 2, "value": "WORLD", "project": {"id": 9}}]},
+            [
+                {"id": 1, "value": "HELLO", "project_id": 8, "project": {"id": 8}},
+                {"id": 2, "value": "WORLD", "project_id": 9, "project": {"id": 9}},
+            ],
+        ),
+    ],
+)
+def test(test_name, transform, body, expected_records):
     extractor = JelloExtractor(transform, decoder, config, kwargs=kwargs)
 
-    records = [{"id": 1}, {"id": 2}]
-    body = {"records": records}
     response = create_response(body)
     actual_records = extractor.extract_records(response)
 
-    assert actual_records == records
+    assert actual_records == expected_records
 
 
 def create_response(body):
     response = requests.Response()
     response._content = json.dumps(body).encode("utf-8")
     return response
-
-
-def test_default():
-    transform = "_{{kwargs['field']}}"
-    extractor = JelloExtractor(transform, decoder, config)
-
-    records = [{"id": 1}, {"id": 2}]
-    response = create_response(records)
-    actual_records = extractor.extract_records(response)
-
-    assert actual_records == records
diff --git a/airbyte-cdk/python/unit_tests/sources/declarative/interpolation/test_interpolated_boolean.py b/airbyte-cdk/python/unit_tests/sources/declarative/interpolation/test_interpolated_boolean.py
index 0fb5bfe64532..eb6b2397083d 100644
--- a/airbyte-cdk/python/unit_tests/sources/declarative/interpolation/test_interpolated_boolean.py
+++ b/airbyte-cdk/python/unit_tests/sources/declarative/interpolation/test_interpolated_boolean.py
@@ -10,6 +10,7 @@
     "string_key": "compare_me",
     "zero_value": 0,
     "empty_array": [],
+    "non_empty_array": [1],
     "empty_dict": {},
     "empty_tuple": (),
 }
@@ -29,6 +30,8 @@
         ("test_empty_dict_is_false", "{{ config['empty_dict'] }}", False),
         ("test_empty_tuple_is_false", "{{ config['empty_tuple'] }}", False),
         ("test_lowercase_false", '{{ "false" }}', False),
+        ("test_value_in_array", "{{ 1 in config['non_empty_array'] }}", True),
+        ("test_value_not_in_array", "{{ 2 in config['non_empty_array'] }}", False),
     ],
 )
 def test_interpolated_boolean(test_name, template, expected_result):