test_kaggle_module_resolver.py

import unittest

import os
import threading
import json

import tensorflow as tf
import tensorflow_hub as hub

from http.server import BaseHTTPRequestHandler, HTTPServer
from test.support.os_helper import EnvironmentVarGuard
from contextlib import contextmanager
from kagglehub.exceptions import BackendError

MOUNT_PATH = "/kaggle/input"

@contextmanager
def create_test_server(handler_class):
    hostname = 'localhost'
    port = 8080
    addr = f"http://{hostname}:{port}"

    # Simulates we are inside a Kaggle environment.
    env = EnvironmentVarGuard()
    env.set('KAGGLE_KERNEL_RUN_TYPE', 'Interactive')
    env.set('KAGGLE_USER_SECRETS_TOKEN', 'foo jwt token')
    env.set('KAGGLE_DATA_PROXY_TOKEN', 'foo proxy token')
    env.set('KAGGLE_DATA_PROXY_URL', addr)

    with env:
        with HTTPServer((hostname, port), handler_class) as test_server:
            threading.Thread(target=test_server.serve_forever).start()

            try:
                yield addr
            finally:
                test_server.shutdown()

class HubHTTPHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        self.send_response(200)
        self.send_header('Content-Type', 'application/gzip')
        self.end_headers()

        with open('/input/tests/data/model.tar.gz', 'rb') as model_archive:
            self.wfile.write(model_archive.read())

class KaggleJwtHandler(BaseHTTPRequestHandler):
    def do_POST(self):
        if self.path.endswith("AttachDatasourceUsingJwtRequest"):
            content_length = int(self.headers["Content-Length"])
            request = json.loads(self.rfile.read(content_length))
            model_ref = request["modelRef"]

            self.send_response(200)
            self.send_header("Content-type", "application/json")
            self.end_headers()

            if model_ref['ModelSlug'] == 'unknown':
                self.wfile.write(bytes(json.dumps({
                    "wasSuccessful": False,
                }), "utf-8"))
                return

            # Load the files
            mount_slug = f"{model_ref['ModelSlug']}/{model_ref['Framework']}/{model_ref['InstanceSlug']}/{model_ref['VersionNumber']}"
            os.makedirs(os.path.dirname(os.path.join(MOUNT_PATH, mount_slug)))
            os.symlink('/input/tests/data/saved_model/', os.path.join(MOUNT_PATH, mount_slug), target_is_directory=True)

            # Return the response
            self.wfile.write(bytes(json.dumps({
                "wasSuccessful": True,
                "result": {
                    "mountSlug": mount_slug,
                },
            }), "utf-8"))
        else:
            self.send_response(404)
            self.wfile.write(bytes(f"Unhandled path: {self.path}", "utf-8"))

class TestKaggleModuleResolver(unittest.TestCase):
    def test_kaggle_resolver_long_url_succeeds(self):
        model_url = "https://kaggle.com/models/foo/foomodule/frameworks/TensorFlow2/variations/barvar/versions/2"
        with create_test_server(KaggleJwtHandler) as addr:
            test_inputs = tf.ones([1,4])
            layer = hub.KerasLayer(model_url)
            self.assertEqual([1, 1], layer(test_inputs).shape)
        # Delete the files that were created in KaggleJwtHandler's do_POST method
        os.unlink(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2"))
        os.rmdir(os.path.dirname(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")))

    def test_kaggle_resolver_short_url_succeeds(self):
        model_url = "https://kaggle.com/models/foo/foomodule/TensorFlow2/barvar/2"
        with create_test_server(KaggleJwtHandler) as addr:
            test_inputs = tf.ones([1,4])
            layer = hub.KerasLayer(model_url)
            self.assertEqual([1, 1], layer(test_inputs).shape)
        # Delete the files that were created in KaggleJwtHandler's do_POST method
        os.unlink(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2"))
        os.rmdir(os.path.dirname(os.path.join(MOUNT_PATH, "foomodule/tensorflow2/barvar/2")))

    def test_kaggle_resolver_not_attached_throws(self):
        with create_test_server(KaggleJwtHandler) as addr:
            with self.assertRaises(BackendError):
                hub.KerasLayer("https://kaggle.com/models/foo/unknown/frameworks/TensorFlow2/variations/barvar/versions/2")

    def test_http_resolver_succeeds(self):
        with create_test_server(HubHTTPHandler) as addr:
            test_inputs = tf.ones([1,4])
            layer = hub.KerasLayer(f'{addr}/model.tar.gz')
            self.assertEqual([1, 1], layer(test_inputs).shape)

    def test_local_path_resolver_succeeds(self):
        test_inputs = tf.ones([1,4])
        layer = hub.KerasLayer('/input/tests/data/saved_model')
        
        self.assertEqual([1, 1], layer(test_inputs).shape)