Typescript match function arguments to key and value of a type map

I have a very simple case, but my experience with Typescript typings is limited and I cannot seem to solve this seemingly simple case.

I have a type map for example like this:

interface KeyValueMap {
  key: 'value';
  foo: 'bar';
}

Now I would like to type the first and second argument of a function to the key and value of the above map

const test = <K extends keyof KeyValueMap>(key: K, value: KeyValueMap[K]) => {
  switch (key) {
    case 'key':
      return value; // currently getting "KeyValueMap[K]" expecting "value"
    case 'foo':
      return value; // currently getting "KeyValueMap[K]" expecting "bar"
  }
};

I tried to search for a similar case, but it seems my Google is broken... So no offense taken when marked as duplicate if such example exists already on Stackoverflow.

UPDATE

After the comment from @ShivamSingla under my question I realize I might have not been clear enough in my question. I am not interested in the return values of the function, but would already like the type to be recognized in the actual function logic (within the switch-case). I will change the example to make it more clear:

interface KeyObjectMap {
  foo: {
    key1: 'value1';
  };
  bar: {
    key2: 'value2';
  };
}

const test = <K extends keyof KeyObjectMap>(key: K, value: KeyObjectMap[K]) => {
  switch (key) {
    case 'foo':
      return value.key1; // property 'key1' does not exist on 'KeyObjectMap[K]'
    case 'bar':
      return value.key2; // property 'key2' does not exist on 'KeyObjectMap[K]'
  }
};

Here you can find a Playground with this example

Since key and value are not associated, value can be re-assigned, before switch case. That's why the error. See this answer for better understanding.

interface KeyObjectMap {
  foo: {
    key1: 'value1';
  };
  bar: {
    key2: 'value2';
  };
}

type K1 = keyof KeyObjectMap

const test = <K extends K1 = K1>(key: K, value: KeyObjectMap[K]) => {
  // error here, 'key2' is missing
  value = {
    key1: 'value1',
  }

  // value is actually intersection, subtype `KeyObjectMap` actually
  value = {
    key1: 'value1',
    key2: 'value2',
  }

  // since `key` and `value` are not associated, value can be re-assigned
  // before switch case. That's why the error

  switch (key) {
    case 'foo':
      return value.key1; // property 'key1' does not exist on 'KeyObjectMap[K]'
    case 'bar':
      return value.key2; // property 'key2' does not exist on 'KeyObjectMap[K]'
  }
};

Playground

Possible solution

interface KeyObjectMap {
  foo: {
    key1: 'value1';
  };
  bar: {
    key2: 'value2';
  };
}

const test = <O extends Partial<KeyObjectMap>>(obj: O) => {
  if (obj.foo) {
    return obj.foo.key1
  }

  if (obj.bar) {
    return obj.bar.key2
  }

  return undefined
};

// calling
const c = test({foo: {key1: 'value1'}})

Playground