Spaces:
Runtime error
Runtime error
File size: 1,991 Bytes
565af7a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import pandas as pd
import pytest
from unittest.mock import patch, MagicMock
from src.vector_database_setup import prepare_movie_descriptions, create_vector_database
# Sample DataFrame to use in both tests
@pytest.fixture
def sample_movies_df():
return pd.DataFrame({
'movieId': [1],
'title': ['The Matrix (1999)'],
'clean_title': ['The Matrix'],
'year': ['1999'],
'genres': [['Action', 'Sci-Fi']],
'avg_rating': [4.5],
'rating_count': [1200]
})
def test_prepare_movie_descriptions(sample_movies_df):
updated_df = prepare_movie_descriptions(sample_movies_df.copy())
assert 'description' in updated_df.columns
assert "The Matrix" in updated_df['description'].iloc[0]
assert "1999" in updated_df['description'].iloc[0]
assert "Action" in updated_df['description'].iloc[0]
@patch("src.vector_database_setup.chromadb.PersistentClient")
@patch("src.vector_database_setup.embedding_functions.SentenceTransformerEmbeddingFunction")
@patch("src.vector_database_setup.SentenceTransformer")
def test_create_vector_database(mock_model, mock_embed_func, mock_client, sample_movies_df):
# Add 'description' column just like prepare_movie_descriptions would
sample_movies_df['description'] = sample_movies_df.apply(
lambda row: f"Title: {row['clean_title']}. "
f"Year: {row['year']}. "
f"Genres: {', '.join(row['genres'])}. "
f"Average Rating: {row['avg_rating']:.1f}/5 from {int(row['rating_count'])} users.",
axis=1
)
# Mock Chroma collection and client
mock_collection = MagicMock()
mock_client.return_value.create_collection.return_value = mock_collection
# Run function
collection = create_vector_database(sample_movies_df)
# Assertions
mock_client.return_value.create_collection.assert_called_once()
mock_collection.add.assert_called_once()
assert collection == mock_collection
|