From 27e2e19805b55224dbbb1a5a4a388bcf427194ed Mon Sep 17 00:00:00 2001
From: Mateusz Kulpa <mateusz.kulpa@outlook.com>
Date: Sun, 24 Mar 2024 14:57:24 +0100
Subject: [PATCH] feat: improve type checking of the default variant identifier
 using `NoInfer` utility type

---
 src/runtime/components/ABTest.vue    |  8 ++++----
 src/runtime/composables/useABTest.ts |  4 ++--
 src/runtime/core/variants.ts         |  2 +-
 src/runtime/types/index.ts           | 13 ++++++++-----
 test/variants.test.ts                | 15 +++++++++++++++
 5 files changed, 30 insertions(+), 12 deletions(-)

diff --git a/src/runtime/components/ABTest.vue b/src/runtime/components/ABTest.vue
index cf5e4df..2cc9720 100644
--- a/src/runtime/components/ABTest.vue
+++ b/src/runtime/components/ABTest.vue
@@ -1,4 +1,4 @@
-<script setup lang="ts" generic="TVariantValue extends JSONValue">
+<script setup lang="ts" generic="TVariantKey extends string, TVariantValue extends JSONValue">
 import { useABTest } from '../composables/useABTest'
 import type { JSONValue, Variant } from '../types'
 import { useSlots } from '#imports'
@@ -6,9 +6,9 @@ import { useSlots } from '#imports'
 const props = withDefaults(
   defineProps<{
     id: string
-    variants: Variant<TVariantValue>[]
+    variants: Variant<TVariantKey, TVariantValue>[]
     enabled?: boolean
-    default?: string
+    default?: NoInfer<TVariantKey>
   }>(),
   {
     enabled: undefined,
@@ -16,7 +16,7 @@ const props = withDefaults(
   }
 )
 
-const { result } = useABTest<TVariantValue>({
+const { result } = useABTest({
   id: props.id,
   variants: props.variants,
   enabled: props.enabled,
diff --git a/src/runtime/composables/useABTest.ts b/src/runtime/composables/useABTest.ts
index 59a5c91..e4d8659 100644
--- a/src/runtime/composables/useABTest.ts
+++ b/src/runtime/composables/useABTest.ts
@@ -4,8 +4,8 @@ import { resolveABTestVariant } from '../core/variants'
 
 const STORAGE_KEY_PREFIX = 'ab-test'
 
-export function useABTest<TVariantValue extends JSONValue>(
-  abTest: ABTest<TVariantValue>
+export function useABTest<TVariantKey extends string, TVariantValue extends JSONValue>(
+  abTest: ABTest<TVariantKey, TVariantValue>
 ): ABTestResult<TVariantValue> {
   const storageId = `${STORAGE_KEY_PREFIX}:${abTest.id}`
 
diff --git a/src/runtime/core/variants.ts b/src/runtime/core/variants.ts
index 8818691..61d5355 100644
--- a/src/runtime/core/variants.ts
+++ b/src/runtime/core/variants.ts
@@ -1,7 +1,7 @@
 import type { ABTest, ABTestResult, JSONValue } from '~/src/runtime/types'
 
 export function resolveABTestVariant<TVariantValue extends JSONValue>(
-  abTest: ABTest<TVariantValue>
+  abTest: ABTest<string, TVariantValue>
 ): ABTestResult<TVariantValue> {
   if (abTest.variants.every(variant => variant.weight === undefined)) {
     return {
diff --git a/src/runtime/types/index.ts b/src/runtime/types/index.ts
index ada7f7f..882d874 100644
--- a/src/runtime/types/index.ts
+++ b/src/runtime/types/index.ts
@@ -14,8 +14,8 @@ export interface JSONObject {
 }
 export type JSONArray = JSONValue[]
 
-export interface Variant<TVariantValue extends JSONValue> {
-  id: string
+export interface Variant<TVariantKey extends string, TVariantValue extends JSONValue> {
+  id: TVariantKey
   value: TVariantValue
   weight?: number
 }
@@ -39,9 +39,12 @@ export type ABTestResult<TVariantValue extends JSONValue> =
   | ABTestResultEnabled<TVariantValue>
   | ABTestResultDisabled<TVariantValue>
 
-export interface ABTest<TVariantValue extends JSONValue = JSONValue> {
+export interface ABTest<
+  TVariantKey extends string = string,
+  TVariantValue extends JSONValue = JSONValue,
+> {
   id: string
-  variants: Variant<TVariantValue>[]
+  variants: Variant<TVariantKey, TVariantValue>[]
   /**
    * Indicates whether the A/B test is enabled or not.
    * By default, the test is considered to be enabled.
@@ -53,5 +56,5 @@ export interface ABTest<TVariantValue extends JSONValue = JSONValue> {
    * The id of the default variant to use when the test is not enabled (`enabled` is `false`).
    * This property is required when `enabled` is `false` to specify which variant should be considered the default.
    */
-  default?: string
+  default?: NoInfer<TVariantKey>
 }
diff --git a/test/variants.test.ts b/test/variants.test.ts
index 50545a0..de1e4ad 100644
--- a/test/variants.test.ts
+++ b/test/variants.test.ts
@@ -58,4 +58,19 @@ describe('useABTest', () => {
 
     expect(result).toEqual({ enabled: false, result: undefined })
   })
+
+  test("doesn't return default value when default is wrong identifier", () => {
+    const result = useABTest({
+      id: 'test-id',
+      variants: [
+        { id: 'a', value: 'a' },
+        { id: 'b', value: 'b' },
+      ],
+      enabled: false,
+      // @ts-expect-error
+      default: 'c',
+    })
+
+    expect(result).toEqual({ enabled: false, result: undefined })
+  })
 })