From 8074a802d6549b32ed60d5ee973c50c397018dac Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Sun, 19 May 2024 00:53:27 +1000
Subject: [PATCH] tests(ui): coverage for validateConnectionTypes

---
 .../util/validateConnectionTypes.test.ts      |  2 +-
 .../store/util/validateConnectionTypes.ts     | 30 +++++++++++--------
 2 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/invokeai/frontend/web/src/features/nodes/store/util/validateConnectionTypes.test.ts b/invokeai/frontend/web/src/features/nodes/store/util/validateConnectionTypes.test.ts
index d953fd973f..10344dd349 100644
--- a/invokeai/frontend/web/src/features/nodes/store/util/validateConnectionTypes.test.ts
+++ b/invokeai/frontend/web/src/features/nodes/store/util/validateConnectionTypes.test.ts
@@ -175,7 +175,7 @@ describe(validateConnectionTypes.name, () => {
       it.each(typePairs)('should accept Collection $t1 to Collection $t2', ({ t1, t2 }: TypePair) => {
         const r = validateConnectionTypes(
           { name: t1, isCollection: true, isCollectionOrScalar: false },
-          { name: t2, isCollection: false, isCollectionOrScalar: false }
+          { name: t2, isCollection: true, isCollectionOrScalar: false }
         );
         expect(r).toBe(true);
       });
diff --git a/invokeai/frontend/web/src/features/nodes/store/util/validateConnectionTypes.ts b/invokeai/frontend/web/src/features/nodes/store/util/validateConnectionTypes.ts
index 092279e315..778b33a7b1 100644
--- a/invokeai/frontend/web/src/features/nodes/store/util/validateConnectionTypes.ts
+++ b/invokeai/frontend/web/src/features/nodes/store/util/validateConnectionTypes.ts
@@ -40,18 +40,25 @@ export const validateConnectionTypes = (sourceType: FieldType, targetType: Field
 
   const isCollectionToGenericCollection = targetType.name === 'CollectionField' && sourceType.isCollection;
 
-  const areBothTypesSingle =
-    !sourceType.isCollection &&
-    !sourceType.isCollectionOrScalar &&
-    !targetType.isCollection &&
-    !targetType.isCollectionOrScalar;
+  const isSourceScalar = !sourceType.isCollection && !sourceType.isCollectionOrScalar;
+  const isTargetScalar = !targetType.isCollection && !targetType.isCollectionOrScalar;
+  const isScalarToScalar = isSourceScalar && isTargetScalar;
+  const isScalarToCollectionOrScalar = isSourceScalar && targetType.isCollectionOrScalar;
+  const isCollectionToCollection = sourceType.isCollection && targetType.isCollection;
+  const isCollectionToCollectionOrScalar = sourceType.isCollection && targetType.isCollectionOrScalar;
+  const isCollectionOrScalarToCollectionOrScalar = sourceType.isCollectionOrScalar && targetType.isCollectionOrScalar;
+  const isPluralityMatch =
+    isScalarToScalar ||
+    isCollectionToCollection ||
+    isCollectionToCollectionOrScalar ||
+    isCollectionOrScalarToCollectionOrScalar ||
+    isScalarToCollectionOrScalar;
 
-  const isIntToFloat = areBothTypesSingle && sourceType.name === 'IntegerField' && targetType.name === 'FloatField';
+  const isIntToFloat = sourceType.name === 'IntegerField' && targetType.name === 'FloatField';
+  const isIntToString = sourceType.name === 'IntegerField' && targetType.name === 'StringField';
+  const isFloatToString = sourceType.name === 'FloatField' && targetType.name === 'StringField';
 
-  const isIntOrFloatToString =
-    areBothTypesSingle &&
-    (sourceType.name === 'IntegerField' || sourceType.name === 'FloatField') &&
-    targetType.name === 'StringField';
+  const isSubTypeMatch = isPluralityMatch && (isIntToFloat || isIntToString || isFloatToString);
 
   const isTargetAnyType = targetType.name === 'AnyField';
 
@@ -62,8 +69,7 @@ export const validateConnectionTypes = (sourceType: FieldType, targetType: Field
     isAnythingToCollectionOrScalarOfSameBaseType ||
     isGenericCollectionToAnyCollectionOrCollectionOrScalar ||
     isCollectionToGenericCollection ||
-    isIntToFloat ||
-    isIntOrFloatToString ||
+    isSubTypeMatch ||
     isTargetAnyType
   );
 };