Test events (#4)

* TABLE_DEF

* test_ReadRecordsRequest

* expect coverage

.github/workflows/python-package.yml

@@ -46,7 +46,7 @@ jobs:
       with:
           coverageFile: coverage.xml
           token: ${{ secrets.GITHUB_TOKEN }}
-          thresholdAll: 0.0
-          thresholdNew: 0.0
-          thresholdModified: 0.0
+          thresholdAll: 0.8
+          thresholdNew: 0.8
+          thresholdModified: 0.8
       if: github.event_name == 'pull_request'

athena_federation/lambda_handler.py

@@ -35,7 +35,6 @@ def process_event(self, event):
         # Look up the request type, call it dynamically, and return the dictionary representation of it.
         # Each model returned implements `as_dict` that returns the info necessary for Athena, including
         # specific PyArrow serialization.
-        print(dir(self))
         request_attr = getattr(self, request_type)
         if request_attr:
             return request_attr().as_dict()

tests/handler.py

@@ -14,8 +14,8 @@
 
 def lambda_handler(event, context):
     # For debugging purposes, we print both the event and the response :)
-    print(json.dumps(event))
+    print("EVENT", json.dumps(event))
     response = example_handler.process_event(event)
-    print(json.dumps(response))
+    print("RESPONSE", json.dumps(response))
 
     return response

tests/test_handler.py

@@ -1,5 +1,7 @@
 from .handler import lambda_handler
 
+TABLE_DEF = {"tableName": "demo", "schemaName": "sampledb"}
+
 
 def test_handler():
     # This is a simple test that checks if the lambda handler can be called
@@ -13,7 +15,7 @@ def test_handler():
     assert True
 
 
-def test_ping():
+def test_PingRequest():
     result = lambda_handler(
         {
             "@type": "PingRequest",
@@ -25,10 +27,11 @@ def test_ping():
         {},
     )
     assert result
+    assert result["@type"] == "PingResponse"
 
 
-def test_schemas():
-    lambda_handler(
+def test_ListSchemasRequest():
+    result = lambda_handler(
         {
             "@type": "ListSchemasRequest",
             "catalogName": "athena_python_sdk",
@@ -37,4 +40,124 @@ def test_schemas():
         },
         {},
     )
-    assert True
+    assert result
+    assert result["@type"] == "ListSchemasResponse"
+
+
+def test_ListTablesRequest():
+    result = lambda_handler(
+        {
+            "@type": "ListTablesRequest",
+            "catalogName": "athena_python_sdk",
+            "schemaName": "sampledb",
+            "tables": ["demo"],
+            "requestType": "LIST_TABLES",
+        },
+        {},
+    )
+    assert result
+    assert result["@type"] == "ListTablesResponse"
+    assert result["requestType"] == "LIST_TABLES"
+    assert "tables" in result
+    assert len(result["tables"]) == 1
+    assert result["tables"][0] == TABLE_DEF
+
+
+def test_GetTableRequest():
+    result = lambda_handler(
+        {
+            "@type": "GetTableRequest",
+            "catalogName": "athena_python_sdk",
+            "tableName": TABLE_DEF,
+        },
+        {},
+    )
+    assert result
+    assert result["@type"] == "GetTableResponse"
+    assert result["requestType"] == "GET_TABLE"
+    schema = result["schema"]
+    assert schema
+    assert "AAAA" in schema["schema"]
+
+
+def test_GetTableLayoutRequest():
+    result = lambda_handler(
+        {
+            "@type": "GetTableLayoutRequest",
+            "catalogName": "athena_python_sdk",
+            "tableName": TABLE_DEF,
+        },
+        {},
+    )
+    assert result
+    assert result["@type"] == "GetTableLayoutResponse"
+    assert result["requestType"] == "GET_TABLE_LAYOUT"
+    assert "aId" in result["partitions"]
+
+
+def test_GetSplitsRequest():
+    result = lambda_handler(
+        {
+            "@type": "GetSplitsRequest",
+            "catalogName": "athena_python_sdk",
+            "tableName": TABLE_DEF,
+        },
+        {},
+    )
+    assert result
+    assert result["@type"] == "GetSplitsResponse"
+    assert result["requestType"] == "GET_SPLITS"
+    splits = result["splits"]
+    assert splits
+    assert len(splits) == 2
+    split = splits[0]
+    print("SPLIT", split)
+    loc = split["spillLocation"]
+    assert loc["bucket"] == "quilt-example"
+    assert "athena-spill" in loc["key"]
+    assert loc["directory"]
+    assert split["properties"] == {"name": "split1", "action": "normal"}
+
+    split1 = splits[1]
+    loc1 = split1["spillLocation"]
+    assert loc1["bucket"] == "quilt-example"
+    assert "athena-spill" in loc1["key"]
+    assert loc1["directory"]
+    assert split1["properties"] == {"name": "split2", "action": "spill"}
+
+
+def test_ReadRecordsRequest():
+    table = lambda_handler(
+        {
+            "@type": "GetTableRequest",
+            "catalogName": "athena_python_sdk",
+            "tableName": TABLE_DEF,
+        },
+        {},
+    )
+    source = lambda_handler(
+        {
+            "@type": "GetSplitsRequest",
+            "catalogName": "athena_python_sdk",
+            "tableName": TABLE_DEF,
+        },
+        {},
+    )
+    assert source
+    source["@type"] = "ReadRecordsRequest"
+    source["tableName"] = TABLE_DEF
+    source["schema"] = table["schema"]
+    source["split"] = source["splits"][0]
+    result = lambda_handler(
+        source,
+        {},
+    )
+    assert result
+    assert result["@type"] == "ReadRecordsResponse"
+    assert result["requestType"] == "READ_RECORDS"
+    assert "records" in result
+    records = result["records"]
+    assert len(records) == 3
+    assert "aId" in records
+    assert "schema" in records
+    assert "records" in records