(t *testing.T)
| 518 | } |
| 519 | |
| 520 | func TestQueryDatabase_Truncation(t *testing.T) { |
| 521 | databases := []map[string]any{ |
| 522 | makeDatabase("instances/prod-pg/databases/employee_db", "instances/prod-pg", "projects/hr-system", "POSTGRES", "ds-admin-1"), |
| 523 | } |
| 524 | // Generate 150 rows. |
| 525 | rows := make([][]any, 150) |
| 526 | for i := range rows { |
| 527 | rows[i] = []any{fmt.Sprintf("%d", i+1)} |
| 528 | } |
| 529 | qr := makeQueryResponse([]string{"id"}, []string{"int4"}, rows, "0.050s") |
| 530 | s := newTestServerWithMock(t, mockQueryServer(databases, qr)) |
| 531 | |
| 532 | result, structured, err := s.handleQueryDatabase(testContext(), nil, QueryInput{ |
| 533 | Database: "employee_db", |
| 534 | Statement: "SELECT id FROM users", |
| 535 | // Default limit is 100. |
| 536 | }) |
| 537 | require.NoError(t, err) |
| 538 | require.False(t, result.IsError) |
| 539 | |
| 540 | output, ok := structured.(*QueryOutput) |
| 541 | require.True(t, ok) |
| 542 | require.Equal(t, 100, output.RowCount) |
| 543 | require.True(t, output.Truncated) |
| 544 | |
| 545 | text := result.Content[0].(*mcpsdk.TextContent).Text |
| 546 | require.Contains(t, text, "Truncated") |
| 547 | } |
| 548 | |
| 549 | func TestQueryDatabase_CustomLimit(t *testing.T) { |
| 550 | databases := []map[string]any{ |
nothing calls this directly
no test coverage detected