diff --git a/athena_federation/lambda_handler.py b/athena_federation/lambda_handler.py index da62450..2373efc 100644 --- a/athena_federation/lambda_handler.py +++ b/athena_federation/lambda_handler.py @@ -27,7 +27,7 @@ def process_event(self, event): # Populate attributes needed to process the event self.event = event self.catalog_name = self.event.get("catalogName") - request_type = self.event.get("@type") or "csv" + request_type = self.event.get("@type") or "PingRequest" # TODO: As the event comes in, populate the database name, table name if provided. # TODO: I can also create new Request types from the event @@ -35,9 +35,10 @@ 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. - response = getattr(self, request_type)() - if response: - return response.as_dict() + print(dir(self)) + request_attr = getattr(self, request_type) + if request_attr: + return request_attr().as_dict() else: print(f"Unknown request type: {request_type}") return {} diff --git a/tests/test_handler.py b/tests/test_handler.py index 8d132e8..eadcf73 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -1,9 +1,40 @@ from .handler import lambda_handler -def test_lambda_handler(): +def test_handler(): # This is a simple test that checks if the lambda handler can be called # without throwing an exception. - lambda_handler({}, {}) + lambda_handler( + { + "queryId": "1681559a-548b-4771-874c-2aa2ea7c39ab", + }, + {}, + ) assert True + +def test_ping(): + result = lambda_handler( + { + "@type": "PingRequest", + "catalogName": "athena_python_sdk", + "queryId": "1681559a-548b-4771-874c-2aa2ea7c39ab", + "sourceType": "athena_python_sdk", + "capabilities": 23, + }, + {}, + ) + assert result + + +def test_schemas(): + lambda_handler( + { + "@type": "ListSchemasRequest", + "catalogName": "athena_python_sdk", + "schemas": ["sampledb"], + "requestType": "LIST_SCHEMAS", + }, + {}, + ) + assert True