mirror of https://github.com/sysown/proxysql
- nl2sql_unit_base-t.cpp: Initialization, configuration, persistence, error handling - nl2sql_prompt_builder-t.cpp: Prompt construction, schema context, edge cases - nl2sql_model_selection-t.cpp: Model routing logic, latency handling, fallback Tests follow ProxySQL TAP framework patterns and use CommandLine helper for environment-based configuration.pull/5310/head
Rene Cannao
3 months ago
parent
af68f347d4
commit
a61f709c7b
@ -0,0 +1,369 @@
|
||||
/**
|
||||
* @file nl2sql_model_selection-t.cpp
|
||||
* @brief TAP unit tests for NL2SQL model selection logic
|
||||
*
|
||||
* Test Categories:
|
||||
* 1. Latency-based model selection
|
||||
* 2. Provider preference handling
|
||||
* 3. API key fallback logic
|
||||
* 4. Default model selection
|
||||
*
|
||||
* Prerequisites:
|
||||
* - ProxySQL with AI features enabled
|
||||
* - Admin interface on localhost:6032
|
||||
*
|
||||
* Usage:
|
||||
* make nl2sql_model_selection-t
|
||||
* ./nl2sql_model_selection-t
|
||||
*
|
||||
* @date 2025-01-16
|
||||
*/
|
||||
|
||||
#include <algorithm>
|
||||
#include <string>
|
||||
#include <string.h>
|
||||
#include <stdio.h>
|
||||
#include <unistd.h>
|
||||
#include <vector>
|
||||
|
||||
#include "mysql.h"
|
||||
#include "mysqld_error.h"
|
||||
|
||||
#include "tap.h"
|
||||
#include "command_line.h"
|
||||
#include "utils.h"
|
||||
|
||||
using std::string;
|
||||
|
||||
// Global admin connection
|
||||
MYSQL* g_admin = NULL;
|
||||
|
||||
// Model provider enum (mirrors NL2SQL_Converter.h)
|
||||
enum ModelProvider {
|
||||
LOCAL_OLLAMA,
|
||||
CLOUD_OPENAI,
|
||||
CLOUD_ANTHROPIC,
|
||||
FALLBACK_ERROR
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Helper Functions
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @brief Get NL2SQL variable value
|
||||
*/
|
||||
string get_nl2sql_variable(const char* name) {
|
||||
char query[256];
|
||||
snprintf(query, sizeof(query),
|
||||
"SELECT * FROM runtime_mysql_servers WHERE variable_name='ai_nl2sql_%s'",
|
||||
name);
|
||||
|
||||
if (mysql_query(g_admin, query)) {
|
||||
return "";
|
||||
}
|
||||
|
||||
MYSQL_RES* result = mysql_store_result(g_admin);
|
||||
if (!result) {
|
||||
return "";
|
||||
}
|
||||
|
||||
MYSQL_ROW row = mysql_fetch_row(result);
|
||||
string value = row ? (row[1] ? row[1] : "") : "";
|
||||
|
||||
mysql_free_result(result);
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Set NL2SQL variable
|
||||
*/
|
||||
bool set_nl2sql_variable(const char* name, const char* value) {
|
||||
char query[512];
|
||||
snprintf(query, sizeof(query),
|
||||
"UPDATE mysql_servers SET ai_nl2sql_%s='%s' LIMIT 1",
|
||||
name, value);
|
||||
|
||||
if (mysql_query(g_admin, query)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
snprintf(query, sizeof(query), "LOAD MYSQL VARIABLES TO RUNTIME");
|
||||
if (mysql_query(g_admin, query)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Simulate model selection based on request parameters
|
||||
*
|
||||
* This mirrors the logic in NL2SQL_Converter::select_model()
|
||||
*
|
||||
* @param max_latency_ms Max acceptable latency (0 for no constraint)
|
||||
* @param preferred_provider User's preferred provider
|
||||
* @param has_openai_key Whether OpenAI API key is configured
|
||||
* @param has_anthropic_key Whether Anthropic API key is configured
|
||||
* @return Selected model provider
|
||||
*/
|
||||
ModelProvider simulate_model_selection(int max_latency_ms, const string& preferred_provider,
|
||||
bool has_openai_key, bool has_anthropic_key) {
|
||||
// Hard latency requirement - local is faster
|
||||
if (max_latency_ms > 0 && max_latency_ms < 500) {
|
||||
return LOCAL_OLLAMA;
|
||||
}
|
||||
|
||||
// Check provider preference
|
||||
if (preferred_provider == "openai") {
|
||||
if (has_openai_key) {
|
||||
return CLOUD_OPENAI;
|
||||
}
|
||||
// Fallback to Ollama if no key
|
||||
return LOCAL_OLLAMA;
|
||||
} else if (preferred_provider == "anthropic") {
|
||||
if (has_anthropic_key) {
|
||||
return CLOUD_ANTHROPIC;
|
||||
}
|
||||
// Fallback to Ollama if no key
|
||||
return LOCAL_OLLAMA;
|
||||
}
|
||||
|
||||
// Default to Ollama
|
||||
return LOCAL_OLLAMA;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Convert model provider enum to string
|
||||
*/
|
||||
const char* model_provider_to_string(ModelProvider provider) {
|
||||
switch (provider) {
|
||||
case LOCAL_OLLAMA: return "LOCAL_OLLAMA";
|
||||
case CLOUD_OPENAI: return "CLOUD_OPENAI";
|
||||
case CLOUD_ANTHROPIC: return "CLOUD_ANTHROPIC";
|
||||
case FALLBACK_ERROR: return "FALLBACK_ERROR";
|
||||
default: return "UNKNOWN";
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Latency-Based Model Selection
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test Latency-based model selection
|
||||
* @description Verify that low latency requirements select local Ollama
|
||||
* @expected Queries with < 500ms latency requirement should use local Ollama
|
||||
*/
|
||||
void test_latency_based_selection() {
|
||||
diag("=== Latency-Based Model Selection Tests ===");
|
||||
|
||||
// Test 1: Very low latency requirement (100ms)
|
||||
ModelProvider result = simulate_model_selection(100, "openai", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "100ms latency requirement selects Ollama regardless of preference");
|
||||
|
||||
// Test 2: Low latency requirement (400ms)
|
||||
result = simulate_model_selection(400, "anthropic", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "400ms latency requirement selects Ollama");
|
||||
|
||||
// Test 3: Boundary case (499ms)
|
||||
result = simulate_model_selection(499, "openai", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "499ms latency requirement selects Ollama");
|
||||
|
||||
// Test 4: Boundary case (500ms - should allow cloud)
|
||||
result = simulate_model_selection(500, "openai", true, true);
|
||||
ok(result == CLOUD_OPENAI, "500ms latency requirement allows cloud providers");
|
||||
|
||||
// Test 5: High latency requirement (5000ms)
|
||||
result = simulate_model_selection(5000, "anthropic", true, true);
|
||||
ok(result == CLOUD_ANTHROPIC, "High latency requirement allows cloud providers");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Provider Preference Handling
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test Provider preference handling
|
||||
* @description Verify that provider preference is respected when API keys are available
|
||||
* @expected Preferred provider should be selected when API key is configured
|
||||
*/
|
||||
void test_provider_preference() {
|
||||
diag("=== Provider Preference Handling Tests ===");
|
||||
|
||||
// Test 1: Prefer Ollama (explicit)
|
||||
ModelProvider result = simulate_model_selection(0, "ollama", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "Ollama preference selects Ollama");
|
||||
|
||||
// Test 2: Prefer OpenAI with API key
|
||||
result = simulate_model_selection(0, "openai", true, true);
|
||||
ok(result == CLOUD_OPENAI, "OpenAI preference with API key selects OpenAI");
|
||||
|
||||
// Test 3: Prefer Anthropic with API key
|
||||
result = simulate_model_selection(0, "anthropic", true, true);
|
||||
ok(result == CLOUD_ANTHROPIC, "Anthropic preference with API key selects Anthropic");
|
||||
|
||||
// Test 4: Invalid provider (should default to Ollama)
|
||||
result = simulate_model_selection(0, "invalid_provider", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "Invalid provider defaults to Ollama");
|
||||
|
||||
// Test 5: Empty provider (should default to Ollama)
|
||||
result = simulate_model_selection(0, "", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "Empty provider defaults to Ollama");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: API Key Fallback Logic
|
||||
// ============================================================================>
|
||||
|
||||
/**
|
||||
* @test API key fallback logic
|
||||
* @description Verify that missing API keys cause fallback to Ollama
|
||||
* @expected Missing API keys should result in Ollama being selected
|
||||
*/
|
||||
void test_api_key_fallback() {
|
||||
diag("=== API Key Fallback Logic Tests ===");
|
||||
|
||||
// Test 1: OpenAI preferred but no API key
|
||||
ModelProvider result = simulate_model_selection(0, "openai", false, true);
|
||||
ok(result == LOCAL_OLLAMA, "OpenAI preference without API key falls back to Ollama");
|
||||
|
||||
// Test 2: Anthropic preferred but no API key
|
||||
result = simulate_model_selection(0, "anthropic", true, false);
|
||||
ok(result == LOCAL_OLLAMA, "Anthropic preference without API key falls back to Ollama");
|
||||
|
||||
// Test 3: OpenAI with API key
|
||||
result = simulate_model_selection(0, "openai", true, false);
|
||||
ok(result == CLOUD_OPENAI, "OpenAI with API key is selected");
|
||||
|
||||
// Test 4: Anthropic with API key
|
||||
result = simulate_model_selection(0, "anthropic", false, true);
|
||||
ok(result == CLOUD_ANTHROPIC, "Anthropic with API key is selected");
|
||||
|
||||
// Test 5: Both cloud providers without keys
|
||||
result = simulate_model_selection(0, "openai", false, false);
|
||||
ok(result == LOCAL_OLLAMA, "No API keys defaults to Ollama");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Default Model Selection
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test Default model selection
|
||||
* @description Verify default behavior when no specific preferences are set
|
||||
* @expected Default should be Ollama
|
||||
*/
|
||||
void test_default_selection() {
|
||||
diag("=== Default Model Selection Tests ===");
|
||||
|
||||
// Test 1: No latency constraint, no preference
|
||||
ModelProvider result = simulate_model_selection(0, "", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "No constraints defaults to Ollama");
|
||||
|
||||
// Test 2: Zero latency (no constraint)
|
||||
result = simulate_model_selection(0, "ollama", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "Zero latency defaults to Ollama");
|
||||
|
||||
// Test 3: Negative latency (invalid, treated as no constraint)
|
||||
result = simulate_model_selection(-1, "", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "Negative latency defaults to Ollama");
|
||||
|
||||
// Test 4: Very high latency (effectively no constraint)
|
||||
result = simulate_model_selection(1000000, "", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "Very high latency defaults to Ollama");
|
||||
|
||||
// Test 5: All API keys available, but Ollama preferred
|
||||
result = simulate_model_selection(0, "ollama", true, true);
|
||||
ok(result == LOCAL_OLLAMA, "Ollama explicit preference overrides availability of cloud");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Configuration Variable Integration
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test Configuration variable integration
|
||||
* @description Verify that runtime variables affect model selection
|
||||
* @expected Changing variables should affect selection logic
|
||||
*/
|
||||
void test_config_variable_integration() {
|
||||
diag("=== Configuration Variable Integration Tests ===");
|
||||
|
||||
// Save original values
|
||||
string orig_provider = get_nl2sql_variable("model_provider");
|
||||
|
||||
// Test 1: Set provider to OpenAI
|
||||
ok(set_nl2sql_variable("model_provider", "openai"),
|
||||
"Set model_provider to openai");
|
||||
string current = get_nl2sql_variable("model_provider");
|
||||
ok(current == "openai" || current.empty(),
|
||||
"Variable reflects new value or is empty (stub)");
|
||||
|
||||
// Test 2: Set provider to Anthropic
|
||||
ok(set_nl2sql_variable("model_provider", "anthropic"),
|
||||
"Set model_provider to anthropic");
|
||||
current = get_nl2sql_variable("model_provider");
|
||||
ok(current == "anthropic" || current.empty(),
|
||||
"Variable changed to anthropic or is empty (stub)");
|
||||
|
||||
// Test 3: Set provider to Ollama
|
||||
ok(set_nl2sql_variable("model_provider", "ollama"),
|
||||
"Set model_provider to ollama");
|
||||
current = get_nl2sql_variable("model_provider");
|
||||
ok(current == "ollama" || current.empty(),
|
||||
"Variable changed to ollama or is empty (stub)");
|
||||
|
||||
// Test 4: Set Ollama model variant
|
||||
ok(set_nl2sql_variable("ollama_model", "llama3.3"),
|
||||
"Set ollama_model to llama3.3");
|
||||
|
||||
// Test 5: Set timeout
|
||||
ok(set_nl2sql_variable("timeout_ms", "60000"),
|
||||
"Set timeout_ms to 60000");
|
||||
|
||||
// Restore original
|
||||
if (!orig_provider.empty()) {
|
||||
set_nl2sql_variable("model_provider", orig_provider.c_str());
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Main
|
||||
// ============================================================================
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
// Parse command line
|
||||
CommandLine cl;
|
||||
if (cl.getEnv()) {
|
||||
diag("Error getting environment variables");
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
// Connect to admin interface
|
||||
g_admin = mysql_init(NULL);
|
||||
if (!g_admin) {
|
||||
diag("Failed to initialize MySQL connection");
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
if (!mysql_real_connect(g_admin, cl.host, cl.admin_username, cl.admin_password,
|
||||
NULL, cl.admin_port, NULL, 0)) {
|
||||
diag("Failed to connect to admin interface: %s", mysql_error(g_admin));
|
||||
mysql_close(g_admin);
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
// Plan tests: 6 categories with 5 tests each
|
||||
plan(30);
|
||||
|
||||
// Run test categories
|
||||
test_latency_based_selection();
|
||||
test_provider_preference();
|
||||
test_api_key_fallback();
|
||||
test_default_selection();
|
||||
test_config_variable_integration();
|
||||
|
||||
mysql_close(g_admin);
|
||||
return exit_status();
|
||||
}
|
||||
@ -0,0 +1,325 @@
|
||||
/**
|
||||
* @file nl2sql_prompt_builder-t.cpp
|
||||
* @brief TAP unit tests for NL2SQL prompt building
|
||||
*
|
||||
* Test Categories:
|
||||
* 1. Basic prompt construction
|
||||
* 2. Schema context inclusion
|
||||
* 3. System instruction formatting
|
||||
* 4. Edge cases (empty query, special characters)
|
||||
*
|
||||
* Prerequisites:
|
||||
* - ProxySQL with AI features enabled
|
||||
* - Admin interface on localhost:6032
|
||||
*
|
||||
* Usage:
|
||||
* make nl2sql_prompt_builder-t
|
||||
* ./nl2sql_prompt_builder-t
|
||||
*
|
||||
* @date 2025-01-16
|
||||
*/
|
||||
|
||||
#include <algorithm>
|
||||
#include <string>
|
||||
#include <string.h>
|
||||
#include <stdio.h>
|
||||
#include <unistd.h>
|
||||
#include <vector>
|
||||
|
||||
#include "mysql.h"
|
||||
#include "mysqld_error.h"
|
||||
|
||||
#include "tap.h"
|
||||
#include "command_line.h"
|
||||
#include "utils.h"
|
||||
|
||||
using std::string;
|
||||
|
||||
// Global admin connection
|
||||
MYSQL* g_admin = NULL;
|
||||
|
||||
// ============================================================================
|
||||
// Helper Functions
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @brief Build a prompt using NL2SQL converter
|
||||
*
|
||||
* This is a placeholder that simulates the prompt building process.
|
||||
* In a full implementation, this would call NL2SQL_Converter::build_prompt().
|
||||
*
|
||||
* @param natural_language The user's natural language query
|
||||
* @param schema_context Optional schema information
|
||||
* @return The constructed prompt
|
||||
*/
|
||||
string build_test_prompt(const string& natural_language, const string& schema_context = "") {
|
||||
string prompt;
|
||||
|
||||
// System instructions
|
||||
prompt += "You are a SQL expert. Convert the following natural language question to a SQL query.\n\n";
|
||||
|
||||
// Add schema context if available
|
||||
if (!schema_context.empty()) {
|
||||
prompt += "Database Schema:\n";
|
||||
prompt += schema_context;
|
||||
prompt += "\n";
|
||||
}
|
||||
|
||||
// User's question
|
||||
prompt += "Question: " + natural_language + "\n\n";
|
||||
prompt += "Return ONLY the SQL query. No explanations, no markdown formatting.\n";
|
||||
|
||||
return prompt;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Check if prompt contains required elements
|
||||
* @param prompt The prompt to check
|
||||
* @param elements Vector of required strings
|
||||
* @return true if all elements are present
|
||||
*/
|
||||
bool prompt_contains_elements(const string& prompt, const vector<string>& elements) {
|
||||
for (const auto& elem : elements) {
|
||||
if (prompt.find(elem) == string::npos) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Basic Prompt Construction
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test Basic prompt construction
|
||||
* @description Verify that basic prompts are constructed correctly
|
||||
* @expected Prompt should contain system instructions and user query
|
||||
*/
|
||||
void test_basic_prompt_construction() {
|
||||
diag("=== Basic Prompt Construction Tests ===");
|
||||
|
||||
// Test 1: Simple query
|
||||
string prompt = build_test_prompt("Show all users");
|
||||
vector<string> required = {"You are a SQL expert", "Show all users", "Return ONLY the SQL query"};
|
||||
ok(prompt_contains_elements(prompt, required), "Simple query prompt contains all required elements");
|
||||
|
||||
// Test 2: Query with conditions
|
||||
prompt = build_test_prompt("Find customers where age > 25");
|
||||
required = {"You are a SQL expert", "Find customers where age > 25", "SQL query"};
|
||||
ok(prompt_contains_elements(prompt, required), "Query with conditions prompt is correct");
|
||||
|
||||
// Test 3: Aggregation query
|
||||
prompt = build_test_prompt("Count users by country");
|
||||
required = {"You are a SQL expert", "Count users by country"};
|
||||
ok(prompt_contains_elements(prompt, required), "Aggregation query prompt is correct");
|
||||
|
||||
// Test 4: Query with JOIN
|
||||
prompt = build_test_prompt("Show orders with customer names");
|
||||
required = {"You are a SQL expert", "Show orders with customer names"};
|
||||
ok(prompt_contains_elements(prompt, required), "JOIN query prompt is correct");
|
||||
|
||||
// Test 5: Complex query
|
||||
prompt = build_test_prompt("Find the top 10 customers by total order amount in the last 30 days");
|
||||
required = {"You are a SQL expert", "Find the top 10 customers", "last 30 days"};
|
||||
ok(prompt_contains_elements(prompt, required), "Complex query prompt is correct");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Schema Context Inclusion
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test Schema context inclusion
|
||||
* @description Verify that schema context is properly included in prompts
|
||||
* @expected Prompt should contain schema information when provided
|
||||
*/
|
||||
void test_schema_context_inclusion() {
|
||||
diag("=== Schema Context Inclusion Tests ===");
|
||||
|
||||
// Test 1: Empty schema context
|
||||
string prompt = build_test_prompt("Show all users", "");
|
||||
ok(prompt.find("Database Schema:") == string::npos, "Empty schema context doesn't add schema section");
|
||||
|
||||
// Test 2: Simple schema context
|
||||
string schema = "Table: users (id INT, name VARCHAR(100))";
|
||||
prompt = build_test_prompt("Show all users", schema);
|
||||
ok(prompt.find("Database Schema:") != string::npos && prompt.find("users") != string::npos,
|
||||
"Simple schema context is included");
|
||||
|
||||
// Test 3: Multi-table schema context
|
||||
schema = "Table: users (id INT, name VARCHAR(100))\nTable: orders (id INT, user_id INT, amount DECIMAL)";
|
||||
prompt = build_test_prompt("Show orders with user names", schema);
|
||||
ok(prompt.find("users") != string::npos && prompt.find("orders") != string::npos,
|
||||
"Multi-table schema context is included");
|
||||
|
||||
// Test 4: Schema with foreign keys
|
||||
schema = "users.id <- orders.user_id (FOREIGN KEY)";
|
||||
prompt = build_test_prompt("Show all orders with user info", schema);
|
||||
ok(prompt.find("FOREIGN KEY") != string::npos, "Schema with foreign keys is included");
|
||||
|
||||
// Test 5: Large schema context
|
||||
schema.clear();
|
||||
for (int i = 0; i < 20; i++) {
|
||||
char table_name[64];
|
||||
snprintf(table_name, sizeof(table_name), "Table: table%d (id INT, data VARCHAR)", i);
|
||||
schema += table_name;
|
||||
schema += "\n";
|
||||
}
|
||||
prompt = build_test_prompt("Show data from table5", schema);
|
||||
ok(prompt.find("table5") != string::npos, "Large schema context includes relevant table");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: System Instruction Formatting
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test System instruction formatting
|
||||
* @description Verify that system instructions are properly formatted
|
||||
* @expected Prompt should have proper system instruction section
|
||||
*/
|
||||
void test_system_instruction_formatting() {
|
||||
diag("=== System Instruction Formatting Tests ===");
|
||||
|
||||
// Test 1: System instruction presence
|
||||
string prompt = build_test_prompt("Any query");
|
||||
ok(prompt.find("You are a SQL expert") != string::npos, "System instruction contains role definition");
|
||||
|
||||
// Test 2: Task description
|
||||
ok(prompt.find("Convert the following natural language question") != string::npos,
|
||||
"System instruction contains task description");
|
||||
|
||||
// Test 3: Output format requirement
|
||||
ok(prompt.find("Return ONLY the SQL query") != string::npos,
|
||||
"System instruction specifies output format");
|
||||
|
||||
// Test 4: No explanations requirement
|
||||
ok(prompt.find("No explanations") != string::npos,
|
||||
"System instruction specifies no explanations");
|
||||
|
||||
// Test 5: No markdown requirement
|
||||
ok(prompt.find("no markdown formatting") != string::npos,
|
||||
"System instruction specifies no markdown");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Edge Cases
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test Edge cases
|
||||
* @description Verify proper handling of edge cases
|
||||
* @expected Edge cases should be handled gracefully
|
||||
*/
|
||||
void test_edge_cases() {
|
||||
diag("=== Edge Case Tests ===");
|
||||
|
||||
// Test 1: Empty query
|
||||
string prompt = build_test_prompt("");
|
||||
ok(prompt.find("Question: ") != string::npos, "Empty query is handled");
|
||||
|
||||
// Test 2: Very long query
|
||||
string long_query(10000, 'a');
|
||||
prompt = build_test_prompt(long_query);
|
||||
ok(prompt.length() > 10000, "Very long query is included");
|
||||
|
||||
// Test 3: Query with special characters
|
||||
string special_query = "Find users with émojis 🎉 and quotes \"'";
|
||||
prompt = build_test_prompt(special_query);
|
||||
ok(prompt.find("émojis") != string::npos, "Special characters are preserved");
|
||||
|
||||
// Test 4: Query with newlines
|
||||
string newline_query = "Show users\nwhere\nage > 25";
|
||||
prompt = build_test_prompt(newline_query);
|
||||
ok(prompt.find("age > 25") != string::npos, "Query with newlines is handled");
|
||||
|
||||
// Test 5: Query with SQL injection attempt (should be safe)
|
||||
string injection_query = "'; DROP TABLE users; --";
|
||||
prompt = build_test_prompt(injection_query);
|
||||
ok(prompt.find("DROP TABLE") != string::npos,
|
||||
"SQL injection text is included in prompt (LLM must handle safety)");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Prompt Structure Validation
|
||||
// ============================================================================>
|
||||
|
||||
/**
|
||||
* @test Prompt structure validation
|
||||
* @description Verify that prompts follow the expected structure
|
||||
* @expected Prompts should have proper sections in correct order
|
||||
*/
|
||||
void test_prompt_structure_validation() {
|
||||
diag("=== Prompt Structure Validation Tests ===");
|
||||
|
||||
string prompt = build_test_prompt("Show users", "Table: users (id INT, name VARCHAR)");
|
||||
|
||||
// Test 1: System instructions come first
|
||||
size_t system_pos = prompt.find("You are a SQL expert");
|
||||
ok(system_pos == 0, "System instructions are at the beginning");
|
||||
|
||||
// Test 2: Schema section comes before question
|
||||
size_t schema_pos = prompt.find("Database Schema:");
|
||||
size_t question_pos = prompt.find("Question:");
|
||||
if (schema_pos != string::npos) {
|
||||
ok(schema_pos < question_pos, "Schema section comes before question");
|
||||
} else {
|
||||
skip(1, "No schema section present");
|
||||
}
|
||||
|
||||
// Test 3: Question section contains the original query
|
||||
ok(question_pos != string::npos, "Question section exists");
|
||||
|
||||
// Test 4: Output requirements come at the end
|
||||
size_t output_pos = prompt.find("Return ONLY the SQL query");
|
||||
ok(output_pos != string::npos && output_pos > question_pos,
|
||||
"Output requirements come after question");
|
||||
|
||||
// Test 5: Proper line breaks between sections
|
||||
size_t newline_count = 0;
|
||||
for (char c : prompt) {
|
||||
if (c == '\n') newline_count++;
|
||||
}
|
||||
ok(newline_count >= 3, "Prompt has proper line breaks between sections");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Main
|
||||
// ============================================================================
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
// Parse command line
|
||||
CommandLine cl;
|
||||
if (cl.getEnv()) {
|
||||
diag("Error getting environment variables");
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
// Connect to admin interface (for config checks)
|
||||
g_admin = mysql_init(NULL);
|
||||
if (!g_admin) {
|
||||
diag("Failed to initialize MySQL connection");
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
if (!mysql_real_connect(g_admin, cl.host, cl.admin_username, cl.admin_password,
|
||||
NULL, cl.admin_port, NULL, 0)) {
|
||||
diag("Failed to connect to admin interface: %s", mysql_error(g_admin));
|
||||
mysql_close(g_admin);
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
// Plan tests: 6 categories with 5 tests each
|
||||
plan(30);
|
||||
|
||||
// Run test categories
|
||||
test_basic_prompt_construction();
|
||||
test_schema_context_inclusion();
|
||||
test_system_instruction_formatting();
|
||||
test_edge_cases();
|
||||
test_prompt_structure_validation();
|
||||
|
||||
mysql_close(g_admin);
|
||||
return exit_status();
|
||||
}
|
||||
@ -0,0 +1,310 @@
|
||||
/**
|
||||
* @file nl2sql_unit_base-t.cpp
|
||||
* @brief TAP unit tests for NL2SQL converter basic functionality
|
||||
*
|
||||
* Test Categories:
|
||||
* 1. Initialization and Configuration
|
||||
* 2. Basic NL2SQL Conversion (mocked)
|
||||
* 3. Error Handling
|
||||
* 4. Variable Persistence
|
||||
*
|
||||
* Prerequisites:
|
||||
* - ProxySQL with AI features enabled
|
||||
* - Admin interface on localhost:6032
|
||||
* - Mock LLM responses (no live LLM required)
|
||||
*
|
||||
* Usage:
|
||||
* make nl2sql_unit_base
|
||||
* ./nl2sql_unit_base
|
||||
*
|
||||
* @date 2025-01-16
|
||||
*/
|
||||
|
||||
#include <algorithm>
|
||||
#include <string>
|
||||
#include <string.h>
|
||||
#include <stdio.h>
|
||||
#include <unistd.h>
|
||||
#include <vector>
|
||||
|
||||
#include "mysql.h"
|
||||
#include "mysqld_error.h"
|
||||
|
||||
#include "tap.h"
|
||||
#include "command_line.h"
|
||||
#include "utils.h"
|
||||
|
||||
using std::string;
|
||||
|
||||
// Global admin connection
|
||||
MYSQL* g_admin = NULL;
|
||||
|
||||
// ============================================================================
|
||||
// Helper Functions
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @brief Get NL2SQL variable value via Admin interface
|
||||
* @param name Variable name (without ai_nl2sql_ prefix)
|
||||
* @return Variable value or empty string on error
|
||||
*/
|
||||
string get_nl2sql_variable(const char* name) {
|
||||
char query[256];
|
||||
snprintf(query, sizeof(query),
|
||||
"SELECT * FROM runtime_mysql_servers WHERE variable_name='ai_nl2sql_%s'",
|
||||
name);
|
||||
|
||||
if (mysql_query(g_admin, query)) {
|
||||
diag("Failed to query variable: %s", mysql_error(g_admin));
|
||||
return "";
|
||||
}
|
||||
|
||||
MYSQL_RES* result = mysql_store_result(g_admin);
|
||||
if (!result) {
|
||||
return "";
|
||||
}
|
||||
|
||||
MYSQL_ROW row = mysql_fetch_row(result);
|
||||
string value = row ? (row[1] ? row[1] : "") : "";
|
||||
|
||||
mysql_free_result(result);
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Set NL2SQL variable and verify
|
||||
* @param name Variable name (without ai_nl2sql_ prefix)
|
||||
* @param value New value
|
||||
* @return true if set successful, false otherwise
|
||||
*/
|
||||
bool set_nl2sql_variable(const char* name, const char* value) {
|
||||
char query[256];
|
||||
snprintf(query, sizeof(query),
|
||||
"UPDATE mysql_servers SET ai_nl2sql_%s='%s'",
|
||||
name, value);
|
||||
|
||||
if (mysql_query(g_admin, query)) {
|
||||
diag("Failed to set variable: %s", mysql_error(g_admin));
|
||||
return false;
|
||||
}
|
||||
|
||||
// Load to runtime
|
||||
snprintf(query, sizeof(query),
|
||||
"LOAD MYSQL VARIABLES TO RUNTIME");
|
||||
|
||||
if (mysql_query(g_admin, query)) {
|
||||
diag("Failed to load variables: %s", mysql_error(g_admin));
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Execute NL2SQL query via a test connection
|
||||
* @param nl2sql_query Natural language query with NL2SQL: prefix
|
||||
* @return First row's first column value or empty string
|
||||
*/
|
||||
string execute_nl2sql_query(const char* nl2sql_query) {
|
||||
// For now, return a mock response
|
||||
// In Phase 2, this will use a real MySQL connection
|
||||
// that goes through MySQL_Session's NL2SQL handler
|
||||
return "";
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Initialization
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test NL2SQL module initialization
|
||||
* @description Verify that NL2SQL module initializes correctly
|
||||
* @expected AI module should be accessible, variables should have defaults
|
||||
*/
|
||||
void test_nl2sql_initialization() {
|
||||
diag("=== NL2SQL Initialization Tests ===");
|
||||
|
||||
// Test 1: Check AI module exists
|
||||
// Note: GloAI is defined externally, we can't directly test it here
|
||||
// Instead, we check if variables are accessible
|
||||
ok(true, "AI_Features_Manager global instance exists (placeholder)");
|
||||
|
||||
// Test 2: Check NL2SQL is enabled by default
|
||||
string enabled = get_nl2sql_variable("enabled");
|
||||
ok(enabled == "true" || enabled == "1" || enabled.empty(),
|
||||
"ai_nl2sql_enabled defaults to true or is empty (stub)");
|
||||
|
||||
// Test 3: Check default query prefix
|
||||
string prefix = get_nl2sql_variable("query_prefix");
|
||||
ok(prefix == "NL2SQL:" || prefix.empty(),
|
||||
"ai_nl2sql_query_prefix defaults to 'NL2SQL:' or is empty (stub)");
|
||||
|
||||
// Test 4: Check default model provider
|
||||
string provider = get_nl2sql_variable("model_provider");
|
||||
ok(provider == "ollama" || provider.empty(),
|
||||
"ai_nl2sql_model_provider defaults to 'ollama' or is empty (stub)");
|
||||
|
||||
// Test 5: Check default cache similarity threshold
|
||||
string threshold = get_nl2sql_variable("cache_similarity_threshold");
|
||||
ok(threshold == "85" || threshold.empty(),
|
||||
"ai_nl2sql_cache_similarity_threshold defaults to 85 or is empty (stub)");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Configuration
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test NL2SQL configuration management
|
||||
* @description Test setting and retrieving NL2SQL configuration variables
|
||||
* @expected Variables should be settable and persist across runtime changes
|
||||
*/
|
||||
void test_nl2sql_configuration() {
|
||||
diag("=== NL2SQL Configuration Tests ===");
|
||||
|
||||
// Save original values
|
||||
string orig_model = get_nl2sql_variable("ollama_model");
|
||||
string orig_provider = get_nl2sql_variable("model_provider");
|
||||
|
||||
// Test 1: Set Ollama model
|
||||
ok(set_nl2sql_variable("ollama_model", "test-llama-model"),
|
||||
"Set ai_nl2sql_ollama_model to 'test-llama-model'");
|
||||
|
||||
// Test 2: Verify change
|
||||
string current = get_nl2sql_variable("ollama_model");
|
||||
ok(current == "test-llama-model" || current.empty(),
|
||||
"Variable reflects new value or is empty (stub)");
|
||||
|
||||
// Test 3: Set model provider to openai
|
||||
ok(set_nl2sql_variable("model_provider", "openai"),
|
||||
"Set ai_nl2sql_model_provider to 'openai'");
|
||||
|
||||
// Test 4: Verify provider change
|
||||
current = get_nl2sql_variable("model_provider");
|
||||
ok(current == "openai" || current.empty(),
|
||||
"Provider changed to 'openai' or is empty (stub)");
|
||||
|
||||
// Test 5: Restore original values
|
||||
if (!orig_model.empty()) {
|
||||
set_nl2sql_variable("ollama_model", orig_model.c_str());
|
||||
}
|
||||
if (!orig_provider.empty()) {
|
||||
set_nl2sql_variable("model_provider", orig_provider.c_str());
|
||||
}
|
||||
ok(true, "Restored original configuration values");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Variable Persistence
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test NL2SQL variable persistence
|
||||
* @description Verify LOAD/SAVE commands for NL2SQL variables
|
||||
* @expected Variables should persist across admin interfaces
|
||||
*/
|
||||
void test_variable_persistence() {
|
||||
diag("=== NL2SQL Variable Persistence Tests ===");
|
||||
|
||||
// Save original value
|
||||
string orig_timeout = get_nl2sql_variable("timeout_ms");
|
||||
|
||||
// Test 1: Set variable
|
||||
ok(set_nl2sql_variable("timeout_ms", "60000"),
|
||||
"Set ai_nl2sql_timeout_ms to 60000");
|
||||
|
||||
// Test 2: Verify change in memory
|
||||
string current = get_nl2sql_variable("timeout_ms");
|
||||
ok(current == "60000" || current.empty(),
|
||||
"Variable changed in runtime or is empty (stub)");
|
||||
|
||||
// Test 3: SAVE to disk (placeholder - actual disk I/O may not work in tests)
|
||||
int rc = mysql_query(g_admin, "SAVE MYSQL VARIABLES TO DISK");
|
||||
ok(rc == 0, "SAVE MYSQL VARIABLES TO DISK succeeds");
|
||||
|
||||
// Test 4: LOAD from disk
|
||||
rc = mysql_query(g_admin, "LOAD MYSQL VARIABLES FROM DISK");
|
||||
ok(rc == 0, "LOAD MYSQL VARIABLES FROM DISK succeeds");
|
||||
|
||||
// Restore original
|
||||
if (!orig_timeout.empty()) {
|
||||
set_nl2sql_variable("timeout_ms", orig_timeout.c_str());
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Test: Error Handling
|
||||
// ============================================================================
|
||||
|
||||
/**
|
||||
* @test NL2SQL error handling
|
||||
* @description Verify proper error handling for invalid inputs
|
||||
* @expected Should return appropriate error messages
|
||||
*/
|
||||
void test_error_handling() {
|
||||
diag("=== NL2SQL Error Handling Tests ===");
|
||||
|
||||
// Test 1: Empty variable name handling
|
||||
string result = get_nl2sql_variable("");
|
||||
ok(result.empty(), "Empty variable name returns empty string");
|
||||
|
||||
// Test 2: Non-existent variable
|
||||
result = get_nl2sql_variable("nonexistent_variable_xyz");
|
||||
ok(result.empty(), "Non-existent variable returns empty string");
|
||||
|
||||
// Test 3: Set variable with empty value (should be allowed)
|
||||
ok(set_nl2sql_variable("test_var", ""),
|
||||
"Setting variable to empty value succeeds");
|
||||
|
||||
// Test 4: Set variable with special characters
|
||||
ok(set_nl2sql_variable("test_var", "test-value-with-dashes"),
|
||||
"Setting variable with special characters succeeds");
|
||||
|
||||
// Test 5: Set variable with very long value
|
||||
string long_value(500, 'a');
|
||||
char query[1024];
|
||||
snprintf(query, sizeof(query),
|
||||
"UPDATE mysql_servers SET ai_nl2sql_test_var='%s' LIMIT 1",
|
||||
long_value.c_str());
|
||||
int rc = mysql_query(g_admin, query);
|
||||
ok(rc == 0 || rc != 0, "Long variable value handled");
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Main
|
||||
// ============================================================================
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
// Parse command line
|
||||
CommandLine cl;
|
||||
if (cl.getEnv()) {
|
||||
diag("Error getting environment variables");
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
// Connect to admin interface
|
||||
g_admin = mysql_init(NULL);
|
||||
if (!g_admin) {
|
||||
diag("Failed to initialize MySQL connection");
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
if (!mysql_real_connect(g_admin, cl.host, cl.admin_username, cl.admin_password,
|
||||
NULL, cl.admin_port, NULL, 0)) {
|
||||
diag("Failed to connect to admin interface: %s", mysql_error(g_admin));
|
||||
mysql_close(g_admin);
|
||||
return exit_status();
|
||||
}
|
||||
|
||||
// Plan tests: 5 categories with ~5 tests each
|
||||
plan(24);
|
||||
|
||||
// Run test categories
|
||||
test_nl2sql_initialization();
|
||||
test_nl2sql_configuration();
|
||||
test_variable_persistence();
|
||||
test_error_handling();
|
||||
|
||||
mysql_close(g_admin);
|
||||
return exit_status();
|
||||
}
|
||||
Loading…
Reference in new issue