[feat] register UNETLoader, UpscaleModelLoader, StylemModelLoader... #5324
Conversation
dosubot
bot
added
the
size:L
🎭 Playwright Test Results
⏰ Started at: 09/04/2025, 01:15:47 AM UTC 🚀 Running Tests
⏱️ Please wait while tests are running across all browsers... |
arjansingh
force-pushed
the
feat/register-more-nodes
branch
from
be8e192 to
d0b6a76
Compare
| let store: ReturnType<typeof useModelToNodeStore> | ||
| let nodeDefStore: ReturnType<typeof useNodeDefStore> |
There was a problem hiding this comment.
Not completely necessary, but you might want to check out the comments I left in here.
I generally prefer to minimize how much mutable state there is ever, but especially in test suites.
The initializing of the stores feels like a pretty important part of the Arrange section for any given test, so it'd be nice to have it directly in the it() block.
| // Create minimal mock for testing - only includes 'name' field since that's | ||
| // the only property ModelNodeProvider constructor uses and tests verify | ||
| const createMockNodeDef = (name: string): ComfyNodeDefImpl => { | ||
| return { name } as ComfyNodeDefImpl |
| const mockNodeNames = [ | ||
| 'CheckpointLoaderSimple', | ||
| 'ImageOnlyCheckpointLoader', | ||
| 'LoraLoader', | ||
| 'LoraLoaderModelOnly', | ||
| 'VAELoader', | ||
| 'ControlNetLoader', | ||
| 'UNETLoader', | ||
| 'UpscaleModelLoader', | ||
| 'StyleModelLoader', | ||
| 'GLIGENLoader' | ||
| ] | ||
|
|
||
| const mockNodeDefs: Record<string, ComfyNodeDefImpl> = Object.fromEntries( | ||
| mockNodeNames.map((name) => [name, createMockNodeDef(name)]) | ||
| ) |
There was a problem hiding this comment.
If this is going to be the same every time, it could be worth extracting as a CONSTANT.
Or if it needs to be regenerated because it is mutable/mutated, an extracted method.
| expect(Object.keys(store.modelToNodeMap)).toContain('checkpoints') | ||
| expect(Object.keys(store.modelToNodeMap)).toContain('unet') |
There was a problem hiding this comment.
https://vitest.dev/api/expect.html#expect-arraycontaining
Might be clearer than having separate expectations, unless these are each separately important to check.
| const provider = store.getNodeProvider('checkpoints') | ||
| expect(provider).toBeDefined() | ||
| // Optional chaining used because getNodeProvider() can return undefined for unregistered types | ||
| expect(provider?.nodeDef.name).toBe('CheckpointLoaderSimple') |
There was a problem hiding this comment.
If one piece is possibly undefined, the chain should usually go to the terminus.
Otherwise you'll get that node isn't present on undefined.
| store = useModelToNodeStore() | ||
| nodeDefStore = useNodeDefStore() |
There was a problem hiding this comment.
Also, why is one store and the other nodeDefStore? Feels like playing favorites 😛
| expect(checkpointProviders.map((p) => p.nodeDef.name)).toContain( | ||
| 'CheckpointLoaderSimple' | ||
| ) | ||
| expect(checkpointProviders.map((p) => p.nodeDef.name)).toContain( | ||
| 'ImageOnlyCheckpointLoader' | ||
| ) | ||
|
|
There was a problem hiding this comment.
For all of this style assertion: https://vitest.dev/api/expect.html#expect-objectcontaining
| ] | ||
|
|
||
| expectedTypes.forEach((modelType) => { | ||
| expect(store.getNodeProvider(modelType)).toBeDefined() |
There was a problem hiding this comment.
When you're looping like this, you usually want to add .soft so that you don't have to keep re-running it if there are multiple violations.
…IGENLoader Also added tests for modelToNodeStore
arjansingh
force-pushed
the
feat/register-more-nodes
branch
from
1e5282c to
68bd05e
Compare
| // Mock nodeDefStore dependency - modelToNodeStore relies on this for registration | ||
| // Most tests expect this to be populated; tests that need empty state can override | ||
| const nodeDefStore = useNodeDefStore() | ||
| nodeDefStore.nodeDefsByName = Object.fromEntries( |
There was a problem hiding this comment.
use vitest mocking
| const modelToNodeStore = useModelToNodeStore() | ||
| modelToNodeStore.registerDefaults() | ||
|
|
||
| EXPECTED_DEFAULT_TYPES.forEach((modelType) => { |
|
Oh I never got a chance to review. I think for the model selector dialog, we will need a mapping in the other direction as well so we will need to add a method in the store for that. |
…5324) * [feat] register UNETLoader, UpscaleModelLoader, StylemModelLoader, GLIGENLoader Also added tests for modelToNodeStore * [fix] code review feedback on tests * [fix] typescript bikeshedding * [fix] remove unnecessary interface mocks
GLIGENLoader
Summary
This is in preparation for getting them to work with our new ModelBrowser component.
Also wrote test suite for the store.
Changes
Nothing should be noticeable to the user.
Review Focus
┆Issue is synchronized with this Notion page by Unito