tests/test_app.py

import json
import os
import pytest
from unittest.mock import patch, MagicMock, mock_open
import pandas as pd
from io import StringIO
from flask import Flask

# Import the Flask app and other necessary modules
from backend.app import app, get_leaderboard, get_chart


@pytest.fixture
def client():
    """Create a test client for the Flask app."""
    app.config['TESTING'] = True
    with app.test_client() as client:
        yield client


def test_index_route(client):
    """Test the index route returns the index.html file."""
    with patch('backend.app.send_from_directory') as mock_send:
        mock_send.return_value = "index content"
        response = client.get('/')
        assert response.status_code == 200
        mock_send.assert_called_once_with(app.static_folder, "index.html")


@patch('backend.app.pd.read_csv')
@patch('backend.app.get_dataset_config')
@patch('backend.app.get_leaderboard')
def test_data_files_benchmark(mock_get_leaderboard, mock_get_config, mock_read_csv, client):
    """Test the data_files route with benchmark dataset type."""
    # Mock the DataFrame and config
    mock_df = pd.DataFrame({'model': ['model1', 'model2'], 'score': [0.9, 0.8]})
    mock_read_csv.return_value = mock_df
    mock_config = {'name': 'test_dataset'}
    mock_get_config.return_value = mock_config
    mock_get_leaderboard.return_value = "leaderboard data"

    # Make the request
    response = client.get('/data/test_dataset?dataset_type=benchmark')

    # Assertions
    mock_read_csv.assert_called_once()
    mock_get_config.assert_called_once_with('test_dataset')
    mock_get_leaderboard.assert_called_once_with(mock_config, mock_df)
    assert response == "leaderboard data"


@patch('backend.app.pd.read_csv')
@patch('backend.app.get_dataset_config')
@patch('backend.app.get_chart')
def test_data_files_attacks_variations(mock_get_chart, mock_get_config, mock_read_csv, client):
    """Test the data_files route with attacks_variations dataset type."""
    # Mock the DataFrame and config
    mock_df = pd.DataFrame({'model': ['model1', 'model2'], 'attack': ['a1', 'a2']})
    mock_read_csv.return_value = mock_df
    mock_config = {'name': 'test_dataset'}
    mock_get_config.return_value = mock_config
    mock_get_chart.return_value = "chart data"

    # Make the request
    response = client.get('/data/test_dataset?dataset_type=attacks_variations')

    # Assertions
    mock_read_csv.assert_called_once()
    mock_get_config.assert_called_once_with('test_dataset')
    mock_get_chart.assert_called_once_with(mock_config, mock_df)
    assert response == "chart data"


def test_data_files_missing_dataset_type(client):
    """Test the data_files route with missing dataset_type."""
    response = client.get('/data/test_dataset')
    assert response.status_code == 400
    assert b"Dataset type not specified" in response.data


@patch('backend.app.pd.read_csv')
def test_data_files_file_not_found(mock_read_csv, client):
    """Test the data_files route when file is not found."""
    mock_read_csv.side_effect = FileNotFoundError()

    response = client.get('/data/test_dataset?dataset_type=benchmark')

    assert response.status_code == 404
    assert b"File not found" in response.data


@patch('backend.app.image_examples_tab')
def test_example_files_image(mock_image_examples, client):
    """Test the example_files route with image type."""
    mock_image_examples.return_value = {"examples": ["image1", "image2"]}

    response = client.get('/examples/image')

    assert response.status_code == 200
    data = json.loads(response.data)
    assert "examples" in data
    mock_image_examples.assert_called_once()


@patch('backend.app.audio_examples_tab')
def test_example_files_audio(mock_audio_examples, client):
    """Test the example_files route with audio type."""
    mock_audio_examples.return_value = {"examples": ["audio1", "audio2"]}

    response = client.get('/examples/audio')

    assert response.status_code == 200
    data = json.loads(response.data)
    assert "examples" in data
    mock_audio_examples.assert_called_once()


@patch('backend.app.video_examples_tab')
def test_example_files_video(mock_video_examples, client):
    """Test the example_files route with video type."""
    mock_video_examples.return_value = {"examples": ["video1", "video2"]}

    response = client.get('/examples/video')

    assert response.status_code == 200
    data = json.loads(response.data)
    assert "examples" in data
    mock_video_examples.assert_called_once()


def test_example_files_invalid_type(client):
    """Test the example_files route with invalid type."""
    response = client.get('/examples/invalid')

    assert response.status_code == 400
    assert b"Invalid example type" in response.data


@patch('backend.app.requests.get')
def test_proxy_valid_url(mock_requests_get, client):
    """Test the proxy route with a valid URL."""
    # Create a mock response
    mock_response = MagicMock()
    mock_response.status_code = 200
    mock_response.content = b"test content"
    mock_response.headers = {"Content-Type": "text/plain"}
    mock_requests_get.return_value = mock_response

    # Mock the ABS_DATASET_DOMAIN constant
    with patch('backend.app.ABS_DATASET_DOMAIN', 'https://example.com'):
        response = client.get('/proxy/https%3A%2F%2Fexample.com%2Ftest.txt')

    assert response.status_code == 200
    assert response.data == b"test content"
    mock_requests_get.assert_called_once()


@patch('backend.app.requests.get')
def test_proxy_invalid_domain(mock_requests_get, client):
    """Test the proxy route with an invalid domain."""
    # Mock the ABS_DATASET_DOMAIN constant
    with patch('backend.app.ABS_DATASET_DOMAIN', 'https://example.com'):
        response = client.get('/proxy/https%3A%2F%2Fmalicious.com%2Ftest.txt')

    assert response.status_code == 403
    data = json.loads(response.data)
    assert "error" in data
    mock_requests_get.assert_not_called()


@patch('backend.app.requests.get')
def test_proxy_request_error(mock_requests_get, client):
    """Test the proxy route when the request fails."""
    mock_requests_get.side_effect = Exception("Connection error")

    # Mock the ABS_DATASET_DOMAIN constant
    with patch('backend.app.ABS_DATASET_DOMAIN', 'https://example.com'):
        response = client.get('/proxy/https%3A%2F%2Fexample.com%2Ftest.txt')

    assert response.status_code == 500
    data = json.loads(response.data)
    assert "error" in data


@patch('backend.app.get_old_format_dataframe')
@patch('backend.app.get_leaderboard_filters')
def test_get_leaderboard(mock_get_filters, mock_get_old_format):
    """Test the get_leaderboard function."""
    # Mock the input data
    df = pd.DataFrame({'model': ['model1', 'model2'], 'score': [0.9, 0.8]})
    config = {
        'first_cols': ['model'],
        'attack_scores': ['score'],
        'categories': ['category1']
    }

    # Mock the function returns
    mock_get_old_format.return_value = df
    mock_get_filters.return_value = (
        {'group1': ['metric1', 'metric2']},
        ['metric1']
    )

    # Call the function
    response = get_leaderboard(config, df)

    # Assertions
    assert response.status_code == 200
    data = json.loads(response.data)
    assert 'groups' in data
    assert 'default_selected_metrics' in data
    assert 'rows' in data
    mock_get_old_format.assert_called_once_with(df, config['first_cols'], config['attack_scores'])
    mock_get_filters.assert_called_once_with(df, config['categories'])


@patch('backend.app.mk_variations')
def test_get_chart(mock_mk_variations):
    """Test the get_chart function."""
    # Mock the input data
    df = pd.DataFrame({'model': ['model1', 'model2'], 'attack': ['a1', 'a2']})
    config = {
        'attacks_with_variations': ['attack1', 'attack2']
    }

    # Mock the function return
    mock_chart_data = {'chart': 'data'}
    mock_mk_variations.return_value = mock_chart_data

    # Call the function
    response = get_chart(config, df)

    # Assertions
    assert response.status_code == 200
    data = json.loads(response.data)
    assert data == mock_chart_data
    mock_mk_variations.assert_called_once_with(df, config['attacks_with_variations'])


# Add a test for the main function
def test_main():
    """Test the main function."""
    with patch('backend.app.app.run') as mock_run:
        # We can't directly test __main__ block, but we can test the app.run call
        # by importing and calling it in a test context
        import backend.app
        if hasattr(backend.app, 'main'):  # If main function exists
            backend.app.main()
            mock_run.assert_called_once_with(host="0.0.0.0", port=7860, debug=True, use_reloader=True)