kftray-commons 0.27.33

use log::error;
use sqlx::{
    Row,
    SqlitePool,
};

use crate::db::get_db_pool;
use crate::models::config_state_model::ConfigState;
use crate::utils::db_mode::{
    DatabaseManager,
    DatabaseMode,
};

pub async fn update_config_state_with_pool(
    config_state: &ConfigState, pool: &SqlitePool,
) -> Result<(), String> {
    let mut conn = pool.acquire().await.map_err(|e| e.to_string())?;
    sqlx::query("UPDATE config_state SET is_running = ?1, process_id = ?2, is_retrying = ?3, retry_count = ?4, last_error = ?5 WHERE config_id = ?6")
        .bind(config_state.is_running)
        .bind(config_state.process_id)
        .bind(config_state.is_retrying)
        .bind(config_state.retry_count)
        .bind(&config_state.last_error)
        .bind(config_state.config_id)
        .execute(&mut *conn)
        .await
        .map_err(|e| e.to_string())?;
    Ok(())
}

pub async fn update_config_state(config_state: &ConfigState) -> Result<(), String> {
    let pool = get_db_pool().await.map_err(|e| e.to_string())?;
    update_config_state_with_pool(config_state, &pool).await
}

pub async fn read_config_states_with_pool(
    pool: &SqlitePool,
) -> Result<Vec<ConfigState>, sqlx::Error> {
    let mut conn = pool.acquire().await.map_err(|e| {
        error!("Failed to acquire database connection: {e}");
        e
    })?;
    let rows = sqlx::query("SELECT id, config_id, is_running, process_id, is_retrying, retry_count, last_error FROM config_state")
        .fetch_all(&mut *conn)
        .await
        .map_err(|e| {
            error!("Failed to fetch config states: {e}");
            e
        })?;

    let config_states = rows
        .into_iter()
        .map(|row| {
            let id: Option<i64> = row.try_get("id").ok();
            let config_id: i64 = row.try_get("config_id").map_err(|e| {
                error!("Failed to get config_id: {e}");
                e
            })?;
            let is_running: bool = row.try_get("is_running").map_err(|e| {
                error!("Failed to get is_running: {e}");
                e
            })?;
            let process_id: Option<u32> = row.try_get("process_id").ok().flatten();
            let is_retrying: bool = row.try_get("is_retrying").ok().unwrap_or(false);
            let retry_count: Option<i32> = row.try_get("retry_count").ok().flatten();
            let last_error: Option<String> = row.try_get("last_error").ok().flatten();
            Ok(ConfigState {
                id,
                config_id,
                is_running,
                process_id,
                is_retrying,
                retry_count,
                last_error,
            })
        })
        .collect::<Result<Vec<_>, sqlx::Error>>()?;

    Ok(config_states)
}

pub async fn read_config_states() -> Result<Vec<ConfigState>, sqlx::Error> {
    let pool = get_db_pool()
        .await
        .map_err(|e| sqlx::Error::Configuration(e.into()))?;
    read_config_states_with_pool(&pool).await
}

pub async fn get_configs_state_with_pool(pool: &SqlitePool) -> Result<Vec<ConfigState>, String> {
    read_config_states_with_pool(pool).await.map_err(|e| {
        error!("Failed to get config states: {e}");
        e.to_string()
    })
}

pub async fn get_configs_state() -> Result<Vec<ConfigState>, String> {
    let pool = get_db_pool().await.map_err(|e| e.to_string())?;
    get_configs_state_with_pool(&pool).await
}

pub async fn update_config_state_with_mode(
    config_state: &ConfigState, mode: DatabaseMode,
) -> Result<(), String> {
    let context = DatabaseManager::get_context(mode).await?;
    update_config_state_with_pool(config_state, &context.pool).await
}

pub async fn read_config_states_with_mode(
    mode: DatabaseMode,
) -> Result<Vec<ConfigState>, sqlx::Error> {
    let context = DatabaseManager::get_context(mode)
        .await
        .map_err(|e| sqlx::Error::Configuration(e.into()))?;
    read_config_states_with_pool(&context.pool).await
}

pub async fn get_configs_state_with_mode(mode: DatabaseMode) -> Result<Vec<ConfigState>, String> {
    let context = DatabaseManager::get_context(mode).await?;
    read_config_states_with_pool(&context.pool)
        .await
        .map_err(|e| {
            error!("Failed to get config states: {e}");
            e.to_string()
        })
}

/// Marks every configuration this process was running as stopped, except the
/// ones in `still_owed`: a configuration whose cleanup did not complete keeps
/// its running state, so the next run's stop-all enumerates and retries it.
pub async fn cleanup_current_process_config_states_with_mode(
    mode: DatabaseMode, still_owed: &[i64],
) -> Result<(), String> {
    let current_process_id = std::process::id();
    let context = DatabaseManager::get_context(mode).await?;
    let mut tx = context.pool.begin().await.map_err(|e| e.to_string())?;

    // The snapshots of the rows about to be marked stopped go with them. Both
    // statements run in one transaction: a failure between them must not
    // leave a row marked running with its snapshot already gone, which would
    // make a later stop fall back to a possibly-edited row.
    let stopping: Vec<i64> = sqlx::query("SELECT config_id FROM config_state WHERE process_id = ?")
        .bind(current_process_id)
        .fetch_all(&mut *tx)
        .await
        .map_err(|e| e.to_string())?
        .into_iter()
        .filter_map(|row| row.try_get::<i64, _>("config_id").ok())
        .filter(|id| !still_owed.contains(id))
        .collect();

    if !stopping.is_empty() {
        let mut delete = sqlx::QueryBuilder::new("DELETE FROM settings WHERE key IN (");
        let mut keys = delete.separated(", ");
        for id in &stopping {
            keys.push_bind(running_snapshot_key(*id, mode));
        }
        delete.push(")");
        delete
            .build()
            .execute(&mut *tx)
            .await
            .map_err(|e| e.to_string())?;
    }

    let mut query = sqlx::QueryBuilder::new(
        "UPDATE config_state SET is_running = false, process_id = NULL WHERE process_id = ",
    );
    query.push_bind(current_process_id);
    if !still_owed.is_empty() {
        query.push(" AND config_id NOT IN (");
        let mut ids = query.separated(", ");
        for id in still_owed {
            ids.push_bind(id);
        }
        query.push(")");
    }
    let affected_rows = query
        .build()
        .execute(&mut *tx)
        .await
        .map_err(|e| e.to_string())?
        .rows_affected();

    tx.commit().await.map_err(|e| e.to_string())?;

    if affected_rows > 0 {
        log::info!(
            "Cleaned up {affected_rows} config states for process {current_process_id} in mode {mode:?}"
        );
    }

    Ok(())
}

/// Key under which the configuration a running forward was started with is
/// kept, next to its running state.
///
/// The row can be edited while the forward runs, so another process that
/// needs to know which address and aliases this one holds cannot read the
/// row for it: the snapshot is what the forward actually owns.
pub fn running_snapshot_key(id: i64, mode: DatabaseMode) -> String {
    format!(
        "running_snapshot:{}:{id}",
        crate::utils::settings::mode_scope(mode)
    )
}

/// Records what a forward that has just registered actually holds. Only the
/// file database is shared between processes; nothing is written for the
/// in-memory one.
pub async fn set_running_snapshot(
    id: i64, config: &crate::models::config_model::Config, mode: DatabaseMode,
) -> Result<(), String> {
    if mode != DatabaseMode::File {
        return Ok(());
    }
    let serialized = serde_json::to_string(config)
        .map_err(|error| format!("Failed to describe the running config {id}: {error}"))?;
    crate::utils::settings::set_setting_with_mode(
        &running_snapshot_key(id, mode),
        &serialized,
        mode,
    )
    .await
    .map_err(|error| format!("Failed to record the running config {id}: {error}"))
}

/// Forgets the snapshot once the forward has stopped.
pub async fn clear_running_snapshot(id: i64, mode: DatabaseMode) {
    if mode != DatabaseMode::File {
        return;
    }
    if let Err(error) =
        crate::utils::settings::delete_setting_with_mode(&running_snapshot_key(id, mode), mode)
            .await
    {
        log::debug!("Failed to clear the running snapshot for config {id}: {error}");
    }
}

/// The configuration a running forward was started with, when a snapshot was
/// recorded for it.
pub async fn running_snapshot(
    id: i64, mode: DatabaseMode,
) -> Option<crate::models::config_model::Config> {
    if mode != DatabaseMode::File {
        return None;
    }
    let stored =
        crate::utils::settings::get_setting_with_mode(&running_snapshot_key(id, mode), mode)
            .await
            .ok()
            .flatten()?;
    serde_json::from_str(&stored).ok()
}

/// Whether a process that recorded itself as running a forward still exists.
///
/// A row whose process is gone is this process's to clean up; one whose
/// process is alive belongs to it, and its listener cannot be stopped from
/// here. Only existence is checked: a reused pid is possible but a process
/// that exited cleanly clears its rows first, so a stale row with a live pid
/// means a crash followed by pid reuse, and leaving it alone is the safe
/// reading.
pub fn process_is_alive(pid: u32) -> bool {
    if pid == 0 {
        return false;
    }
    if pid == std::process::id() {
        return true;
    }
    #[cfg(unix)]
    {
        let Ok(pid) = libc::pid_t::try_from(pid) else {
            return false;
        };
        // Signal 0 checks for existence without delivering anything. EPERM
        // means the process exists but belongs to another user.
        // SAFETY: signal 0 delivers nothing; `pid` is a validated, non-zero
        // process id.
        let result = unsafe { libc::kill(pid, 0) };
        result == 0 || std::io::Error::last_os_error().raw_os_error() == Some(libc::EPERM)
    }
    #[cfg(windows)]
    {
        use windows::Win32::Foundation::{
            CloseHandle,
            ERROR_ACCESS_DENIED,
            WAIT_TIMEOUT,
        };
        use windows::Win32::System::Threading::{
            OpenProcess,
            PROCESS_QUERY_LIMITED_INFORMATION,
            PROCESS_SYNCHRONIZE,
            WaitForSingleObject,
        };
        match unsafe {
            OpenProcess(
                PROCESS_QUERY_LIMITED_INFORMATION | PROCESS_SYNCHRONIZE,
                false,
                pid,
            )
        } {
            Ok(handle) => {
                let running = unsafe { WaitForSingleObject(handle, 0) } == WAIT_TIMEOUT;
                let _ = unsafe { CloseHandle(handle) };
                running
            }
            // Access denied still means the process exists. `Error::code`
            // wraps the Win32 status as an HRESULT, so it must be compared
            // against `HRESULT::from_win32(ERROR_ACCESS_DENIED.0)`, not the
            // unrelated COM `E_ACCESSDENIED` constant.
            Err(error) => error.code() == windows::core::HRESULT::from_win32(ERROR_ACCESS_DENIED.0),
        }
    }
    #[cfg(not(any(unix, windows)))]
    {
        true
    }
}

#[cfg(test)]
mod tests {
    use sqlx::SqlitePool;

    use super::*;
    use crate::db::create_db_table;
    use crate::models::config_model::Config;
    use crate::utils::config;

    async fn setup_test_db() -> SqlitePool {
        let pool = SqlitePool::connect("sqlite::memory:")
            .await
            .expect("Failed to connect to in-memory database");
        create_db_table(&pool)
            .await
            .expect("Failed to create tables");
        crate::utils::migration::migrate_configs(Some(&pool))
            .await
            .expect("Failed to run migrations");
        pool
    }

    #[tokio::test]
    async fn test_read_initial_config_state() {
        let pool = setup_test_db().await;
        let config_data = Config {
            service: Some("state-test-1".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();
        let configs = config::read_configs_with_pool(&pool).await.unwrap();
        let config_id = configs.first().unwrap().id.unwrap();
        let states = read_config_states_with_pool(&pool).await.unwrap();
        assert_eq!(states.len(), 1);
        let state = &states[0];
        assert_eq!(state.config_id, config_id);
        assert!(!state.is_running);
        assert!(state.id.is_some());
    }

    #[tokio::test]
    async fn test_update_and_read_config_state() {
        let pool = setup_test_db().await;
        let config_data = Config {
            service: Some("state-test-2".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();
        let configs = config::read_configs_with_pool(&pool).await.unwrap();
        let config_id = configs.first().unwrap().id.unwrap();
        let initial_states = read_config_states_with_pool(&pool).await.unwrap();
        let initial_state = initial_states
            .iter()
            .find(|s| s.config_id == config_id)
            .unwrap();
        assert!(!initial_state.is_running);
        let state_to_update = ConfigState {
            id: initial_state.id,
            config_id,
            is_running: true,
            process_id: Some(1234),
            ..Default::default()
        };
        update_config_state_with_pool(&state_to_update, &pool)
            .await
            .unwrap();
        let updated_states = read_config_states_with_pool(&pool).await.unwrap();
        assert_eq!(updated_states.len(), 1);
        let updated_state = updated_states
            .iter()
            .find(|s| s.config_id == config_id)
            .unwrap();
        assert_eq!(updated_state.config_id, config_id);
        assert!(updated_state.is_running);
        assert_eq!(updated_state.id, initial_state.id);
    }

    #[tokio::test]
    async fn test_read_multiple_config_states() {
        let pool = setup_test_db().await;
        let config1 = Config {
            service: Some("state-test-3".to_string()),
            ..Config::default()
        };
        let config2 = Config {
            service: Some("state-test-4".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config1.clone(), &pool)
            .await
            .unwrap();
        config::insert_config_with_pool(config2.clone(), &pool)
            .await
            .unwrap();
        let configs = config::read_configs_with_pool(&pool).await.unwrap();
        let config1_id = configs
            .iter()
            .find(|c| c.service == config1.service)
            .unwrap()
            .id
            .unwrap();
        let config2_id = configs
            .iter()
            .find(|c| c.service == config2.service)
            .unwrap()
            .id
            .unwrap();
        let state_to_update = ConfigState {
            id: None,
            config_id: config1_id,
            is_running: true,
            process_id: Some(1234),
            ..Default::default()
        };
        update_config_state_with_pool(&state_to_update, &pool)
            .await
            .unwrap();
        let states = read_config_states_with_pool(&pool).await.unwrap();
        assert_eq!(states.len(), 2);
        let state1 = states.iter().find(|s| s.config_id == config1_id).unwrap();
        let state2 = states.iter().find(|s| s.config_id == config2_id).unwrap();
        assert!(state1.is_running);
        assert!(!state2.is_running);
    }

    #[tokio::test]
    async fn test_get_configs_state_wrapper() {
        let pool = setup_test_db().await;
        let config_data = Config {
            service: Some("state-test-wrapper".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();
        let result = read_config_states_with_pool(&pool)
            .await
            .map_err(|e| e.to_string());
        assert!(result.is_ok());
        assert_eq!(result.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_update_config_state_public_function() {
        let pool = setup_test_db().await;
        let config_data = Config {
            service: Some("state-test-5".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();
        let configs = config::read_configs_with_pool(&pool).await.unwrap();
        let config_id = configs.first().unwrap().id.unwrap();

        let state_to_update = ConfigState {
            id: None,
            config_id,
            is_running: true,
            process_id: Some(1234),
            ..Default::default()
        };

        tokio::task::yield_now().await;

        let result = update_config_state_with_pool(&state_to_update, &pool).await;
        assert!(result.is_ok());

        let states = read_config_states_with_pool(&pool).await.unwrap();
        let updated_state = states.iter().find(|s| s.config_id == config_id).unwrap();
        assert!(updated_state.is_running);
    }

    #[tokio::test]
    async fn test_read_config_states_public_function() {
        let pool = setup_test_db().await;
        let config_data = Config {
            service: Some("state-test-6".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();

        let states_result = read_config_states_with_pool(&pool).await;
        assert!(states_result.is_ok());
        assert_eq!(states_result.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_get_configs_state_function() {
        let pool = setup_test_db().await;
        let config_data = Config {
            service: Some("state-test-7".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();

        let states_result = read_config_states_with_pool(&pool)
            .await
            .map_err(|e| e.to_string());
        assert!(states_result.is_ok());
        let states = states_result.unwrap();
        assert_eq!(states.len(), 1);
        assert_eq!(states[0].config_id, 1);
    }

    #[tokio::test]
    async fn test_error_handling_in_update_config_state() {
        let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
        pool.close().await;

        let config_state = ConfigState {
            id: Some(1),
            config_id: 1,
            is_running: true,
            process_id: Some(1234),
            ..Default::default()
        };

        let result = update_config_state_with_pool(&config_state, &pool).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_error_handling_in_read_config_states() {
        let pool = SqlitePool::connect("sqlite::memory:").await.unwrap();
        pool.close().await;

        let result = read_config_states_with_pool(&pool).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_row_processing_in_read_config_states() {
        let pool = setup_test_db().await;

        let config1 = Config {
            service: Some("row-test-1".to_string()),
            ..Config::default()
        };
        let config2 = Config {
            service: Some("row-test-2".to_string()),
            ..Config::default()
        };

        config::insert_config_with_pool(config1.clone(), &pool)
            .await
            .unwrap();
        config::insert_config_with_pool(config2.clone(), &pool)
            .await
            .unwrap();

        let configs = config::read_configs_with_pool(&pool).await.unwrap();
        let config2_id = configs
            .iter()
            .find(|c| c.service == Some("row-test-2".to_string()))
            .unwrap()
            .id
            .unwrap();

        let state = ConfigState {
            id: None,
            config_id: config2_id,
            is_running: true,
            process_id: Some(1234),
            ..Default::default()
        };

        update_config_state_with_pool(&state, &pool).await.unwrap();

        let states = read_config_states_with_pool(&pool).await.unwrap();
        assert_eq!(states.len(), 2);

        let state1 = states.iter().find(|s| s.config_id == 1).unwrap();
        let state2 = states.iter().find(|s| s.config_id == config2_id).unwrap();

        assert!(!state1.is_running);
        assert!(state2.is_running);
    }

    #[tokio::test]
    async fn test_update_config_state_with_direct_pool() {
        let pool = setup_test_db().await;

        let config_data = Config {
            service: Some("public-test-1".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();

        let configs = config::read_configs_with_pool(&pool).await.unwrap();
        let config_id = configs.first().unwrap().id.unwrap();

        let initial_states = read_config_states_with_pool(&pool).await.unwrap();
        assert_eq!(initial_states.len(), 1);
        assert!(!initial_states[0].is_running);

        let _conn = pool.acquire().await.unwrap();

        let state = ConfigState {
            id: initial_states[0].id,
            config_id,
            is_running: true,
            process_id: Some(1234),
            ..Default::default()
        };

        let result = update_config_state_with_pool(&state, &pool).await;
        assert!(result.is_ok());

        let updated_states = read_config_states_with_pool(&pool).await.unwrap();
        assert_eq!(updated_states.len(), 1);
        assert!(updated_states[0].is_running);
    }

    #[tokio::test]
    async fn test_read_config_states_with_direct_pool() {
        let pool = setup_test_db().await;

        let config_data = Config {
            service: Some("public-read-test".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();

        let direct_result = read_config_states_with_pool(&pool).await;
        assert!(direct_result.is_ok());
        assert_eq!(direct_result.unwrap().len(), 1);
    }

    #[tokio::test]
    async fn test_get_configs_state_public_wrapper() {
        let pool = setup_test_db().await;

        let config_data = Config {
            service: Some("get-wrapper-test".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();

        let direct_states = read_config_states_with_pool(&pool).await.unwrap();
        assert_eq!(direct_states.len(), 1);
        let err_result = get_configs_state().await;
        assert!(err_result.is_ok() || err_result.is_err());
    }

    #[tokio::test]
    async fn test_error_handling_row_processing_config_id() {
        let pool = setup_test_db().await;

        let config_data = Config {
            service: Some("error-test".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();

        let states = read_config_states_with_pool(&pool).await.unwrap();
        assert_eq!(states.len(), 1);
        assert_eq!(states[0].config_id, 1);
    }

    #[tokio::test]
    async fn test_config_state_with_mode_memory() {
        let _lock = crate::test_utils::MEMORY_MODE_TEST_MUTEX.lock().await;

        let _ = config::delete_all_configs_with_mode(DatabaseMode::Memory).await;

        let config_data = Config {
            service: Some("memory-state-test".to_string()),
            ..Config::default()
        };

        config::insert_config_with_mode(config_data, DatabaseMode::Memory)
            .await
            .unwrap();

        let configs = config::read_configs_with_mode(DatabaseMode::Memory)
            .await
            .unwrap();
        let config_id = configs[0].id.unwrap();

        let states = read_config_states_with_mode(DatabaseMode::Memory)
            .await
            .unwrap();
        assert_eq!(states.len(), 1);
        assert_eq!(states[0].config_id, config_id);
        assert!(!states[0].is_running);

        let state_update = ConfigState {
            id: None,
            config_id,
            is_running: true,
            process_id: Some(1234),
            ..Default::default()
        };

        update_config_state_with_mode(&state_update, DatabaseMode::Memory)
            .await
            .unwrap();

        let updated_states = read_config_states_with_mode(DatabaseMode::Memory)
            .await
            .unwrap();
        assert_eq!(updated_states.len(), 1);
        assert!(updated_states[0].is_running);
    }

    #[tokio::test]
    async fn test_get_configs_state_with_mode_memory() {
        let _lock = crate::test_utils::MEMORY_MODE_TEST_MUTEX.lock().await;

        let _ = config::delete_all_configs_with_mode(DatabaseMode::Memory).await;

        let config_data = Config {
            service: Some("get-state-memory-test".to_string()),
            ..Config::default()
        };

        config::insert_config_with_mode(config_data, DatabaseMode::Memory)
            .await
            .unwrap();

        let states = get_configs_state_with_mode(DatabaseMode::Memory)
            .await
            .unwrap();
        assert_eq!(states.len(), 1);
        assert!(!states[0].is_running);
    }

    #[test]
    fn test_config_state_default_values() {
        let state = ConfigState::default();
        assert_eq!(state.id, None);
        assert_eq!(state.config_id, 0);
        assert!(!state.is_running);
        assert_eq!(state.process_id, None);
        assert!(!state.is_retrying);
        assert_eq!(state.retry_count, None);
        assert_eq!(state.last_error, None);
    }

    #[test]
    fn test_config_state_backward_compat_json() {
        // Old JSON without new fields (is_retrying, retry_count, last_error)
        let old_json = r#"{"id": 1, "config_id": 42, "is_running": true, "process_id": null}"#;
        let state: ConfigState = serde_json::from_str(old_json).unwrap();
        assert_eq!(state.id, Some(1));
        assert_eq!(state.config_id, 42);
        assert!(state.is_running);
        assert_eq!(state.process_id, None);
        // New fields should have defaults
        assert!(!state.is_retrying);
        assert_eq!(state.retry_count, None);
        assert_eq!(state.last_error, None);
    }

    #[test]
    fn test_config_state_json_serialization_with_new_fields() {
        let state = ConfigState {
            id: Some(1),
            config_id: 42,
            is_running: true,
            process_id: Some(1234),
            is_retrying: true,
            retry_count: Some(3),
            last_error: Some("Connection timeout".to_string()),
        };
        let json = serde_json::to_string(&state).unwrap();
        let deserialized: ConfigState = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.id, Some(1));
        assert_eq!(deserialized.config_id, 42);
        assert!(deserialized.is_running);
        assert_eq!(deserialized.process_id, Some(1234));
        assert!(deserialized.is_retrying);
        assert_eq!(deserialized.retry_count, Some(3));
        assert_eq!(
            deserialized.last_error,
            Some("Connection timeout".to_string())
        );
    }

    #[tokio::test]
    async fn test_config_state_db_round_trip_with_retry_fields() {
        let pool = setup_test_db().await;
        let config_data = Config {
            service: Some("retry-test".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();
        let configs = config::read_configs_with_pool(&pool).await.unwrap();
        let config_id = configs.first().unwrap().id.unwrap();
        let initial_states = read_config_states_with_pool(&pool).await.unwrap();
        let initial_state = initial_states
            .iter()
            .find(|s| s.config_id == config_id)
            .unwrap();
        // Update with retry fields
        let state_to_update = ConfigState {
            id: initial_state.id,
            config_id,
            is_running: true,
            process_id: Some(5678),
            is_retrying: true,
            retry_count: Some(3),
            last_error: Some("Pod restart detected".to_string()),
        };
        update_config_state_with_pool(&state_to_update, &pool)
            .await
            .unwrap();
        // Read back and verify all fields persisted
        let updated_states = read_config_states_with_pool(&pool).await.unwrap();
        let updated_state = updated_states
            .iter()
            .find(|s| s.config_id == config_id)
            .unwrap();
        assert_eq!(updated_state.config_id, config_id);
        assert!(updated_state.is_running);
        assert_eq!(updated_state.process_id, Some(5678));
        assert!(updated_state.is_retrying);
        assert_eq!(updated_state.retry_count, Some(3));
        assert_eq!(
            updated_state.last_error,
            Some("Pod restart detected".to_string())
        );
    }

    #[tokio::test]
    async fn test_config_state_db_round_trip_retry_count_none() {
        let pool = setup_test_db().await;
        let config_data = Config {
            service: Some("retry-none-test".to_string()),
            ..Config::default()
        };
        config::insert_config_with_pool(config_data.clone(), &pool)
            .await
            .unwrap();
        let configs = config::read_configs_with_pool(&pool).await.unwrap();
        let config_id = configs.first().unwrap().id.unwrap();
        let initial_states = read_config_states_with_pool(&pool).await.unwrap();
        let initial_state = initial_states
            .iter()
            .find(|s| s.config_id == config_id)
            .unwrap();
        // Update with is_retrying=true but retry_count=None
        let state_to_update = ConfigState {
            id: initial_state.id,
            config_id,
            is_running: false,
            process_id: None,
            is_retrying: true,
            retry_count: None,
            last_error: None,
        };
        update_config_state_with_pool(&state_to_update, &pool)
            .await
            .unwrap();
        // Read back and verify
        let updated_states = read_config_states_with_pool(&pool).await.unwrap();
        let updated_state = updated_states
            .iter()
            .find(|s| s.config_id == config_id)
            .unwrap();
        assert!(updated_state.is_retrying);
        assert_eq!(updated_state.retry_count, None);
        assert_eq!(updated_state.last_error, None);
    }
}