Improve compilation cache tests

[ayaka14732]

This PR improves the tests for the compilation cache.

Background

The JAX compilation cache operates through a set of functions defined in jax/_src/compilation_cache.py. These functions implement the operations of the compilation cache. A global variable, _cache, is defined in this file as an instance of the abstract class CacheInterface. This instance can be either GFileCache or the new LRUCache added in #21394.

There are three levels of abstraction involved:

1. Compilation cache operations
2. The CacheInterface abstract class
3. Implementations of the CacheInterface abstract class

The Problem

The current compilation_cache_test.py file is intended to test the compilation cache operations but relies heavily on the implementation details of the underlying cache. For instance, it assumes that the number of files in the cache directory matches the number of items in the cache.

Why is This a Problem?

With the intended introduction of the LRU cache eviction policy for the JAX compilation cache, we now require two separate files for each cache item: one for data and one for metadata. Consequently, the assumption that one cache item corresponds to one file in the cache directory is no longer valid.

The Solution

This PR addresses the problem by implementing two functions, count_cache_items() and clear_cache(), at each level of abstraction. This ensures that the tests for the compilation cache do not depend on the specific implementation details of the underlying cache.

[hawkinsp]

Please squash your commits.