tests/test_roomcategory.py

#!/usr/bin/env python3

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


@pytest.fixture
def classifier():
    model_path = "travelgate/room_category-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', 'expected_response': 'comfort'},
        {'input': 'full size bed 1 big bed', 'expected_response': 'default'},
        {'input': '1 king bed suite nonsmoking', 'expected_response': 'default'},
        {'input': 'junior suite 1 king bed nonsmoking', 'expected_response': 'junior'},
        {'input': 'superior sea view', 'expected_response': 'superior'},
        {'input': 'family room superior', 'expected_response': 'superior'},
        {'input': '1 king standard', 'expected_response': 'standard'},
        {'input': 'superior king balcon vista al mar', 'expected_response': 'superior'},
        {'input': 'standard double room king bed accessible', 'expected_response': 'standard'},  # noqa: E501
        {'input': 'sagamore garden view suite standard', 'expected_response': 'standard'},  # noqa: E501
        {'input': 'standard 1 king bed courtyard view non smoking', 'expected_response': 'standard'},  # noqa: E501
        {'input': 'suite double for 2 adults 0 children and 0 infants', 'expected_response': 'default'},  # noqa: E501
        {'input': 'one bedroom superior king suite with ocean view non smoking', 'expected_response': 'superior'},  # noqa: E501
        {'input': 'omnia deluxe triple for 2 adults 0 children and 0 infants', 'expected_response': 'deluxe'},  # noqa: E501
        {'input': 'king premium room incl evening tasting welcome gift comp wifi 27 30 sqm espresso fridge bathrobe', 'expected_response': 'premium'},  # noqa: E501
        {'input': 'standard twin double room single use', 'expected_response': 'standard'},  # noqa: E501
        {'input': 'habitacion 2 camas deluxe', 'expected_response': 'deluxe'},
        {'input': 'standard 2 double beds courtyard view non smoking', 'expected_response': 'standard'},  # noqa: E501
        {'input': 'habitacion estandar con bano', 'expected_response': 'standard'},
        {'input': 'superior room 1 king superior room 1 king cupola or courtyard view french style 36sqm 385sq', 'expected_response': 'superior'},  # noqa: E501
        {'input': 'habitacion estandar matrimonial adaptada discapacitados', 'expected_response': 'standard'},  # noqa: E501
        {'input': 'privilege twin for 2 adults 0 children and 0 infants', 'expected_response': 'privilege'},  # noqa: E501
        {'input': 'classic single for 1 adults 0 children and 0 infants', 'expected_response': 'classic'},  # noqa: E501
        {'input': 'deluxe room double for 2 adults 0 children and 0 infants', 'expected_response': 'deluxe'},  # noqa: E501
        {'input': 'gallery two queens', 'expected_response': 'default'},
        {'input': 'premier palace double room', 'expected_response': 'premier'},
        {'input': 'double single use deluxe', 'expected_response': 'deluxe'},
        {'input': 'double room queen bed superior', 'expected_response': 'superior'},
        {'input': 'double guest standard room', 'expected_response': 'standard'},
        {'input': '2 queen beds disney view non smoking', 'expected_response': 'default'},  # noqa: E501
        {'input': 'standard single room for 1 adults 0 children and 0 infants', 'expected_response': 'standard'},  # noqa: E501
        {'input': 'twin premium room incl evening tasting welcome gift comp wifi 28 sqm espresso fridge bathrobe', 'expected_response': 'premium'},  # noqa: E501
        {'input': 'superior quadruple room', 'expected_response': 'superior'},
        {'input': 'superior one bedroom apartment x2013 2 adults', 'expected_response': 'superior'}  # noqa: E501
    ]

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

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