test handler

athena_federation/lambda_handler.py

@@ -27,15 +27,20 @@ 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")
+        request_type = self.event.get("@type") or "csv"
 
         # 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
 
         # 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.
-        return getattr(self, request_type)().as_dict()
+        response = getattr(self, request_type)()
+        if response:
+            return response.as_dict()
+        else:
+            print(f"Unknown request type: {request_type}")
+            return {}
 
     def PingRequest(self) -> models.PingResponse:
         return models.PingResponse(

athena_federation/example/handler.py → tests/handler.py

@@ -2,10 +2,10 @@
 import os
 
 from athena_federation.lambda_handler import AthenaLambdaHandler
-from sample_data_source import SampleDataSource
+from .sample_data_source import SampleDataSource
 
 # This needs to be a valid bucket that the Lambda function role has access to
-spill_bucket = os.environ["TARGET_BUCKET"]
+spill_bucket = os.environ.get("SPILL_BUCKET", "quilt-example")
 
 example_handler = AthenaLambdaHandler(
     data_source=SampleDataSource(), spill_bucket=spill_bucket

tests/test_handler.py

@@ -0,0 +1,9 @@
+from .handler import lambda_handler
+
+
+def test_lambda_handler():
+    # This is a simple test that checks if the lambda handler can be called
+    # without throwing an exception.
+    lambda_handler({}, {})
+    assert True
+