How to key off a parameter to a stubbed method using Mockito

Greetings.

I am mocking a search engine for testing in my web app. This search engine returns xml documents with different schemas. The schema depends on a parameter known as a collection set. Returning different schemas based on collection sets is the part that's difficult to mock, because specifying the collection set is basically a setup method, and a void one at that. This search engine is an external jar file so I can't modify the API. I have to work with what's been provided. Here's an example:

Engine engine = factory.getEngine();
Search search = engine.getSearch();
search.addCollectionSet(someCollectionSet);
SearchResult result = search.getSearchResult();
Document[] documents = result.getAllDocuments();

Then for each document, I can get the xml by calling:

document.getDocumentText();

When I'm using my mock objects, getDocumentText() returns an xml string, created by a generator, that conforms to the schema. What I want to do is use a different type of generator depending on which collection set was provided in step 3 in the first code snippet above. I've been trying to do something like this:

    doAnswer(new Answer() {
        Object answer(InvocationOnMock invocation) {
            if (args == "foo") {
                SearchResult result = getMockSearchResult();
                when(search.getSearchResult()).thenReturn(result);
            }
        }
    }).when(search.addCollectionSet(anyString()));

But this results in lots of red highlighting :)

Basically, my goal is to key off of addCollectionSet(someCollectionSet) so that when it's called, I can do some kind of switch off of the parameter and ensure that a different generator is used. Does anyone know how I can accomplish something like this? Or is there maybe some form of Dependency Injection that could be used to conditionally wire up my generator?

Thanks!

Update

I've changed my factory object so that it never returns the engine, but rather, the Search and Find objects from that engine, so now I can do something like this:

Search search = factory.getSearch(collectionSet);

So what I'd like to do is something like this:

when(factory.getSearch(anyString()).thenAnswer(new Answer() {
    Object answer(InvocationOnMock invocation) {
        switch(args[0]) {
            case fooSet: return fooSearch; break;
            case barSet: return barSearch; break;   

In other words, I still want to key off the string that was passed in to getSearch in a switch statement. Admittedly, I could开发者_如何学Python do something more like felix has suggested below, but I'd rather have all my cases wrapped in a switch. Can someone provide an example of how this could be done? Thanks!

Update

I've seen that you can capture the arguments that are passed into mocked calls, but these captured arguments are used for later assertions. I haven't seen a way that I can key off these arguments so that a call to my mock will return different values depending on the arguments. It seems like there has to be a way to do this, I just don't have enough experience with Mockito to figure it out. But surely someone does!

I would recommend wrapping the call to the legacy code into your own object. So you end up with your own method along these lines:

class SearchEngineWrapper {
  public String getSearchResult(String collection){
    Engine engine = factory.getEngine();
    Search search = engine.getSearch();
    search.addCollectionSet(someCollectionSet);
    SearchResult result = search.getSearchResult();
    ...
    return document.getDocumentText();
  }
}

Now you can mock out this method. The method also nicely documents your intent. Also you could test the actual implementation in an integration test.

when(searchEngineWrapper.getSearchResult("abc").thenReturn("foo");
when(searchEngineWrapper.getSearchResult("xyz").thenReturn("bar");