Mock a specific ngxs store select with ngxs

In the component there are two ngxs selectors:

@Component({
  selector: 'some-component',
  templateUrl: './some-component.html',
  styleUrls: ['./some-component.scss'],
  changeDetection: ChangeDetectionStrategy.OnPush,
})
export class OverviewComponent implements OnInit {

  @Select(State.foo) foo$: Observable<StateModel['foo']>;
  @Select(State.bar) bar$: Observable<StateModel['bar']>;

  ngOnInit(): void {
   
    combineLatest([foo$, bar$])
       .subscribe(([foo, bar) => {
         this.bar = bar;
       });
  }

  // left out for brevity

Now in the test, the following is possible:

describe('SomeComponent', () => {
  beforeEach(async () => {

    stub.state = {}

    stub.Store = {
      select: jest.fn(),
    }

    await TestBed.configureTestingModule({
      declarations: [SomeComponent],
      imports: [NgxsModule.forRoot([]),
       providers: [{ provide: Store, useValue: stub.Store }]
    }).compileComponents();
  });

  it('does something', () => {
    jest.spyOn(stub.Store, 'select')
      .mockReturnValueOnce(of(false)); // <-- this is the @Select foo$ call
      .mockReturnValueOnce(of(true)); // <-- this is the @Select bar$ call

    expect(component.bar).toBe(true);
  });
});

But it's not really clear that the second .mockReturnValueOnce is for bar$. Now it seems fine, but if I expand my spec and want different outcomes for the selectors than it becomes unclear really fast.

Is there a way to specify which select is being mocked?

Even better, you can use the selector apparently.

So if you have this selector:

export class MySelectors {
  @Selector([State])
  static foo(state: State): string {
    return state.foo;
  }
}

Then you can do:

jest.spyOn(stub.Store, 'select').mockImplementation(selector => {
  if (selector === MySelectors.foo) {
    return of('blaat');
  }
});

And even better, you can overwrite that return value in a spec if you define a setter function:

describe('SomeComponent', () => {
  let fixture: ComponentFixture<SomeComponent>;
  let component: SomeComponent;
  let stub;
  let setupFunction: (
    someValue: string,
  ) => void;

  beforeEach(async () => {
    stub = {
      Store: {
        select: jest.fn(),
      }
    }

    await TestBed.configureTestingModule({
      imports: [ ... ],
      declarations: [ ... ],
      providers: [
        { provide: Store, useValue: stub.Store },
      ]
    }).compileComponents();
    
    fixture = TestBed.createComponent(SomeComponent);
    component = fixture.componentInstance;

    setupStoreSelectors = (someValue) => {
      jest.spyOn(stub.Store, 'select')
        .mockImplementation(selector => {
          if (selector === MySelectors.foo) {
            return of(someValue);
          }
        });
    };
  });

  it('should do something', () => {

    setupStoreSelectors('yay!);
    fixture.detectChanges();

    // Now the value is 'yay!', which you can test here..
  });
});