def setUp(self):
        super(TestCallbacks, self).setUp()

        import chroma_core.lib.storage_plugin.manager

        self.orig_manager = chroma_core.lib.storage_plugin.manager.storage_plugin_manager
        chroma_core.lib.storage_plugin.manager.storage_plugin_manager = (
            chroma_core.lib.storage_plugin.manager.StoragePluginManager()
        )

        from chroma_core.lib.storage_plugin.manager import storage_plugin_manager

        storage_plugin_manager._load_plugin(sys.modules[__name__], "test_mod", TestPlugin)

        from chroma_core.models import StorageResourceRecord

        resource_class, resource_class_id = storage_plugin_manager.get_plugin_resource_class(
            "test_mod", "TestGlobalResource"
        )
        record, created = StorageResourceRecord.get_or_create_root(resource_class, resource_class_id, {"name": "test1"})

        from chroma_core.lib.storage_plugin.query import ResourceQuery

        scannable_record = StorageResourceRecord.objects.get()
        self.scannable_resource = ResourceQuery().get_resource(scannable_record)
        self.scannable_global_id = scannable_record.pk

        self.resource_manager = mock.Mock(_sessions={})
        self.plugin = TestPlugin(self.resource_manager, self.scannable_global_id)
        self.resource_manager._sessions[self.scannable_global_id] = PluginSession(
            self.plugin, self.scannable_global_id, 0
        )
Example #2
0
    def setUp(self):
        super(TestAddRemove, self).setUp()

        import chroma_core.lib.storage_plugin.manager
        self.orig_manager = chroma_core.lib.storage_plugin.manager.storage_plugin_manager
        chroma_core.lib.storage_plugin.manager.storage_plugin_manager = chroma_core.lib.storage_plugin.manager.StoragePluginManager(
        )

        from chroma_core.lib.storage_plugin.manager import storage_plugin_manager
        storage_plugin_manager._load_plugin(sys.modules[__name__], 'test_mod',
                                            TestPlugin)

        from chroma_core.models import StorageResourceRecord
        resource_class, resource_class_id = storage_plugin_manager.get_plugin_resource_class(
            'test_mod', 'TestGlobalResource')
        record, created = StorageResourceRecord.get_or_create_root(
            resource_class, resource_class_id, {'name': 'test1'})

        from chroma_core.lib.storage_plugin.query import ResourceQuery

        scannable_record = StorageResourceRecord.objects.get()
        self.scannable_resource = ResourceQuery().get_resource(
            scannable_record)
        self.scannable_global_id = scannable_record.pk