One more attempt with Claude's help to close connections

tests/unit/utils/inference/test_inference_store.py

@@ -89,6 +89,8 @@ async def test_inference_store_pagination_basic():
         assert result3.data[0].id == "zebra-task"
         assert result3.has_more is False
 
+        await store.sql_store.close()
+
 
 async def test_inference_store_pagination_ascending():
     """Test pagination with ascending order."""
@@ -126,6 +128,8 @@ async def test_inference_store_pagination_ascending():
         assert result2.data[0].id == "charlie-task"
         assert result2.has_more is True
 
+        await store.sql_store.close()
+
 
 async def test_inference_store_pagination_with_model_filter():
     """Test pagination combined with model filtering."""
@@ -166,6 +170,8 @@ async def test_inference_store_pagination_with_model_filter():
         assert result2.data[0].model == "model-a"
         assert result2.has_more is False
 
+        await store.sql_store.close()
+
 
 async def test_inference_store_pagination_invalid_after():
     """Test error handling for invalid 'after' parameter."""
@@ -178,6 +184,8 @@ async def test_inference_store_pagination_invalid_after():
         with pytest.raises(ValueError, match="Record with id='non-existent' not found in table 'chat_completions'"):
             await store.list_chat_completions(after="non-existent", limit=2)
 
+        await store.sql_store.close()
+
 
 async def test_inference_store_pagination_no_limit():
     """Test pagination behavior when no limit is specified."""
@@ -208,3 +216,5 @@ async def test_inference_store_pagination_no_limit():
         assert result.data[0].id == "beta-second"  # Most recent first
         assert result.data[1].id == "omega-first"
         assert result.has_more is False
+
+        await store.sql_store.close()

tests/unit/utils/responses/test_responses_store.py

@@ -98,6 +98,8 @@ async def test_responses_store_pagination_basic():
         assert result3.data[0].id == "zebra-resp"
         assert result3.has_more is False
 
+        await store.sql_store.close()
+
 
 async def test_responses_store_pagination_ascending():
     """Test pagination with ascending order."""
@@ -136,6 +138,8 @@ async def test_responses_store_pagination_ascending():
         assert result2.data[0].id == "charlie-resp"
         assert result2.has_more is True
 
+        await store.sql_store.close()
+
 
 async def test_responses_store_pagination_with_model_filter():
     """Test pagination combined with model filtering."""
@@ -177,6 +181,8 @@ async def test_responses_store_pagination_with_model_filter():
         assert result2.data[0].model == "model-a"
         assert result2.has_more is False
 
+        await store.sql_store.close()
+
 
 async def test_responses_store_pagination_invalid_after():
     """Test error handling for invalid 'after' parameter."""
@@ -189,6 +195,8 @@ async def test_responses_store_pagination_invalid_after():
         with pytest.raises(ValueError, match="Record with id.*'non-existent' not found in table 'openai_responses'"):
             await store.list_responses(after="non-existent", limit=2)
 
+        await store.sql_store.close()
+
 
 async def test_responses_store_pagination_no_limit():
     """Test pagination behavior when no limit is specified."""
@@ -221,6 +229,8 @@ async def test_responses_store_pagination_no_limit():
         assert result.data[1].id == "omega-resp"
         assert result.has_more is False
 
+        await store.sql_store.close()
+
 
 async def test_responses_store_get_response_object():
     """Test retrieving a single response object."""
@@ -249,6 +259,8 @@ async def test_responses_store_get_response_object():
         with pytest.raises(ValueError, match="Response with id non-existent not found"):
             await store.get_response_object("non-existent")
 
+        await store.sql_store.close()
+
 
 async def test_responses_store_input_items_pagination():
     """Test pagination functionality for input items."""
@@ -330,6 +342,8 @@ async def test_responses_store_input_items_pagination():
         with pytest.raises(ValueError, match="Cannot specify both 'before' and 'after' parameters"):
             await store.list_response_input_items("test-resp", before="some-id", after="other-id")
 
+        await store.sql_store.close()
+
 
 async def test_responses_store_input_items_before_pagination():
     """Test before pagination functionality for input items."""
@@ -390,3 +404,5 @@ async def test_responses_store_input_items_before_pagination():
             ValueError, match="Input item with id 'non-existent' not found for response 'test-resp-before'"
         ):
             await store.list_response_input_items("test-resp-before", before="non-existent")
+
+        await store.sql_store.close()

tests/unit/utils/sqlstore/test_sqlstore.py

@@ -64,6 +64,9 @@ async def test_sqlite_sqlstore():
         assert result.data == [{"id": 12, "name": "test12"}]
         assert result.has_more is False
 
+        # cleanup
+        await sqlstore.close()
+
 
 async def test_sqlstore_pagination_basic():
     """Test basic pagination functionality at the SQL store level."""
@@ -128,6 +131,8 @@ async def test_sqlstore_pagination_basic():
         assert result3.data[0]["id"] == "zebra"
         assert result3.has_more is False
 
+        await store.close()
+
 
 async def test_sqlstore_pagination_with_filter():
     """Test pagination with WHERE conditions."""
@@ -180,6 +185,8 @@ async def test_sqlstore_pagination_with_filter():
         assert result2.data[0]["id"] == "xyz"
         assert result2.has_more is False
 
+        await store.close()
+
 
 async def test_sqlstore_pagination_ascending_order():
     """Test pagination with ascending order."""
@@ -228,6 +235,8 @@ async def test_sqlstore_pagination_ascending_order():
         assert result2.data[0]["id"] == "alpha"
         assert result2.has_more is True
 
+        await store.close()
+
 
 async def test_sqlstore_pagination_multi_column_ordering_error():
     """Test that multi-column ordering raises an error when using cursor pagination."""
@@ -265,6 +274,8 @@ async def test_sqlstore_pagination_multi_column_ordering_error():
         assert len(result.data) == 1
         assert result.data[0]["id"] == "task1"
 
+        await store.close()
+
 
 async def test_sqlstore_pagination_cursor_requires_order_by():
     """Test that cursor pagination requires order_by parameter."""
@@ -282,6 +293,8 @@ async def test_sqlstore_pagination_cursor_requires_order_by():
                 cursor=("id", "task1"),
             )
 
+        await store.close()
+
 
 async def test_sqlstore_pagination_error_handling():
     """Test error handling for invalid columns and cursor IDs."""
@@ -414,6 +427,8 @@ async def test_where_operator_edge_cases():
         with pytest.raises(ValueError, match="Unsupported operator"):
             await store.fetch_all("events", where={"ts": {"!=": base}})
 
+        await store.close()
+
 
 async def test_sqlstore_pagination_custom_key_column():
     """Test pagination with custom primary key column (not 'id')."""
@@ -463,3 +478,5 @@ async def test_sqlstore_pagination_custom_key_column():
         assert len(result2.data) == 1
         assert result2.data[0]["uuid"] == "uuid-alpha"
         assert result2.has_more is False
+
+        await store.close()

tests/unit/utils/test_authorized_sqlstore.py

@@ -77,6 +77,8 @@ async def test_authorized_fetch_with_where_sql_access_control(mock_get_authentic
         assert row is not None
         assert row["title"] == "User Document"
 
+        await base_sqlstore.close()
+
 
 @patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
 async def test_sql_policy_consistency(mock_get_authenticated_user):
@@ -163,6 +165,8 @@ async def test_sql_policy_consistency(mock_get_authenticated_user):
                 f"Difference: SQL only: {sql_ids - policy_ids}, Policy only: {policy_ids - sql_ids}"
             )
 
+        await base_sqlstore.close()
+
 
 @patch("llama_stack.providers.utils.sqlstore.authorized_sqlstore.get_authenticated_user")
 async def test_authorized_store_user_attribute_capture(mock_get_authenticated_user):
@@ -211,3 +215,5 @@ async def test_authorized_store_user_attribute_capture(mock_get_authenticated_us
         # Third item should have null attributes (no authenticated user)
         assert result.data[2]["id"] == "item3"
         assert result.data[2]["access_attributes"] is None
+
+        await base_sqlstore.close()