mocking angular $element in jasmine

I need to test a directive controller with $element inside. So I've got a function like this one:

function func($event) {
        $element.find('#la-la-la').focus();
}

and render it in the test:

template = $compile(element)($scope);
$scope.$apply();

controller = element.controller('myDirective');

And what I'm trying to do is to test this function inside that controller for the directive.

describe('func method', function testFunc() {
    it('should focus on element', function checkFocusing() {
        controller.func();
        expect(focusSpy).toHaveBeenCalled();
    });
});

Where "focusSpy" is a spy inside mocked $element service. But it seems like if I use $provide.service('$element', ...) it is not found by the test. Injecting it to $scope.$element before compiling doesn't work either. Thank you!