move chromadb imports and pysqlite hotswap to fix tests

This commit is contained in:
Jason Hunter 2024-06-21 23:51:57 -04:00
parent 35e8bcc68d
commit 6ef58d6fc2
4 changed files with 20 additions and 22 deletions

View File

@ -1,12 +1,9 @@
import faulthandler
import sys
import threading
from flask import cli
# Hotsawp the sqlite3 module for Chroma compatibility
__import__("pysqlite3")
sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
from frigate.app import FrigateApp
faulthandler.enable()
@ -15,8 +12,6 @@ threading.current_thread().name = "frigate"
cli.show_server_banner = lambda *x: None
if __name__ == "__main__":
from frigate.app import FrigateApp
frigate_app = FrigateApp()
frigate_app.start()

View File

@ -54,10 +54,6 @@ def manage_embeddings(config: FrigateConfig) -> None:
models = [Event]
db.bind(models)
# Hotsawp the sqlite3 module for Chroma compatibility
__import__("pysqlite3")
sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
embeddings = Embeddings()
# Check if we need to re-index events

View File

@ -3,19 +3,32 @@
import base64
import io
import logging
import sys
import time
import numpy as np
from chromadb import Collection
from chromadb import HttpClient as ChromaClient
from chromadb.config import Settings
from PIL import Image
from playhouse.shortcuts import model_to_dict
from frigate.models import Event
from .functions.clip import ClipEmbedding
from .functions.minilm_l6_v2 import MiniLMEmbedding
# Hotsawp the sqlite3 module for Chroma compatibility
try:
from chromadb import Collection
from chromadb import HttpClient as ChromaClient
from chromadb.config import Settings
from .functions.clip import ClipEmbedding
from .functions.minilm_l6_v2 import MiniLMEmbedding
except RuntimeError:
__import__("pysqlite3")
sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
from chromadb import Collection
from chromadb import HttpClient as ChromaClient
from chromadb.config import Settings
from .functions.clip import ClipEmbedding
from .functions.minilm_l6_v2 import MiniLMEmbedding
logger = logging.getLogger(__name__)

View File

@ -2,7 +2,6 @@ import datetime
import json
import logging
import os
import sys
import unittest
from unittest.mock import Mock
@ -11,18 +10,13 @@ from playhouse.shortcuts import model_to_dict
from playhouse.sqlite_ext import SqliteExtDatabase
from playhouse.sqliteq import SqliteQueueDatabase
from frigate.api.app import create_app
from frigate.config import FrigateConfig
from frigate.models import Event, Recordings
from frigate.plus import PlusApi
from frigate.stats.emitter import StatsEmitter
from frigate.test.const import TEST_DB, TEST_DB_CLEANUPS
# Have to replace pysqlite3 before importing create_app
__import__("pysqlite3")
sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")
from frigate.api.app import create_app # noqa: E402
class TestHttp(unittest.TestCase):
def setUp(self):