tests/test_roomenvironment.py

#!/usr/bin/env python3

import pytest
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline


@pytest.fixture
def classifier():
    model_path = "travelgate/room_environment-classifier"
    try:
        model = AutoModelForSequenceClassification.from_pretrained(model_path)
        tokenizer = AutoTokenizer.from_pretrained(model_path)
        classification = pipeline(
            "text-classification",
            model=model,
            tokenizer=tokenizer,
            framework="pt",
            device=0 if torch.cuda.is_available() else -1,
        )
        print("Modelo y tokenizer cargados exitosamente.")
    except Exception as e:
        print(f"Error al cargar el modelo o tokenizer: {e}")
        raise e

    return classification


def test_model_predictions(classifier):
    test_data = [
        {'input': 'comfort double with shared bed', 'expected_response': 'shared bed'},
        {'input': 'habitacion con banno compartido', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'habitacion con bano compartido', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'full size bed 1 big bed with shared bathroom', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': '1 king bed suite nonsmoking shared bathroom', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'junior suite 1 king bed nonsmoking shared bed', 'expected_response': 'shared bed'},  # noqa: E501
        {'input': 'superior sea view connected room', 'expected_response': 'connected'},
        {'input': 'twin room with two shared beds', 'expected_response': 'shared bed'},
        {'input': 'double room with connected rooms available', 'expected_response': 'connected'},  # noqa: E501
        {'input': 'two queen beds with shared bathroom access', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'suite with king bed and connected room option', 'expected_response': 'connected'},  # noqa: E501
        {'input': 'shared dormitory room mixed gender', 'expected_response': 'shared room'},  # noqa: E501
        {'input': 'private room with shared bathroom facilities', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'family suite with connected rooms', 'expected_response': 'connected'},  # noqa: E501
        {'input': 'bunk bed in mixed shared room', 'expected_response': 'shared room'},
        {'input': 'deluxe queen room with two queen beds shared bathroom', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'hostel dorm bed shared room', 'expected_response': 'shared room'},
        {'input': 'economy double room with shared bathroom outside the room', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'luxury suite with sea view and private connected rooms', 'expected_response': 'connected'},  # noqa: E501
        {'input': 'single bed in co-ed shared hostel room', 'expected_response': 'shared room'},  # noqa: E501
        {'input': 'executive suite with optional connecting room', 'expected_response': 'connected'},  # noqa: E501
        {'input': 'standard twin room shared beds', 'expected_response': 'shared bed'},
        {'input': 'cozy single room with access to shared bathroom', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'spacious family room with interconnecting doors', 'expected_response': 'connected'},  # noqa: E501
        {'input': 'shared female dormitory for 4 guests', 'expected_response': 'shared room'},  # noqa: E501
        {'input': 'luxury double room with shared bed arrangements', 'expected_response': 'shared bed'},  # noqa: E501
        {'input': 'master suite with a shared bedroom balcony and sea view', 'expected_response': 'shared room'},  # noqa: E501
        {'input': 'budget twin room with shared bath', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'two-bedroom apartment with connecting options', 'expected_response': 'connected'},  # noqa: E501
        {'input': 'en-suite room with shared bathroom', 'expected_response': 'shared bathroom'},  # noqa: E501
        {'input': 'single guest room with shared bed', 'expected_response': 'shared bed'},  # noqa: E501
        {'input': 'shared room in a modern loft city center', 'expected_response': 'shared room'},  # noqa: E501
        {'input': 'double bed in a shared room apartment', 'expected_response': 'shared room'},  # noqa: E501
        {'input': 'studio with shared bath', 'expected_response': 'shared bathroom'},
        {'input': 'penthouse suite connected rooms available upon request', 'expected_response': 'connected'},  # noqa: E501
        {'input': "backpacker's special one bed in a shared room of six", 'expected_response': 'shared room'}  # noqa: E501
    ]

    for test_case in test_data:
        description = test_case["input"]
        expected_environment = test_case["expected_response"]
        result = classifier(description)[0]
        predicted_environment = result["label"]

        assert (
            predicted_environment == expected_environment
        ), f"Incorrect prediction for '{description}': expected '{expected_environment}', obtained '{predicted_environment}'"  # noqa: E501