From 65adfd8ef718797bb1ea433ecbb3e2e89022dd6a Mon Sep 17 00:00:00 2001 From: "Dr. Ernie Prabhakar" Date: Mon, 12 Aug 2024 23:05:19 -0400 Subject: [PATCH] Test events (#4) * TABLE_DEF * test_ReadRecordsRequest * expect coverage --- .github/workflows/python-package.yml | 6 +- athena_federation/lambda_handler.py | 1 - tests/handler.py | 4 +- tests/test_handler.py | 131 ++++++++++++++++++++++++++- 4 files changed, 132 insertions(+), 10 deletions(-) diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index a53cf5d..7805da9 100644 --- a/.github/workflows/python-package.yml +++ b/.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' \ No newline at end of file diff --git a/athena_federation/lambda_handler.py b/athena_federation/lambda_handler.py index 2373efc..212a9a2 100644 --- a/athena_federation/lambda_handler.py +++ b/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() diff --git a/tests/handler.py b/tests/handler.py index 45a4f2a..860529e 100644 --- a/tests/handler.py +++ b/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 diff --git a/tests/test_handler.py b/tests/test_handler.py index eadcf73..64bc325 100644 --- a/tests/test_handler.py +++ b/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