Assistant/cmd/cognitive-worker/model_qualification_test.go

package main

import (
	"context"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"git.cin.su/m/Assistant/adapters/models"
	"git.cin.su/m/Assistant/contracts/capability"
	"git.cin.su/m/Assistant/contracts/cognitive"
)

type qualificationProvider struct{}

func (qualificationProvider) Probe(context.Context) models.ProviderProbe {
	return models.ProviderProbe{Health: capability.HealthAvailable}
}

func (qualificationProvider) Discover(context.Context) ([]models.DiscoveredModel, models.ProviderProbe) {
	return []models.DiscoveredModel{{ProviderModelID: "model-a"}, {ProviderModelID: "model-b"}}, models.ProviderProbe{Health: capability.HealthAvailable}
}
func (qualificationProvider) Advise(_ context.Context, request cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error) {
	if err := models.ValidateAdvisoryRequest(request); err != nil {
		return cognitive.AdvisoryResponse{}, err
	}
	return cognitive.AdvisoryResponse{Version: request.Version, RequestID: request.RequestID, AdvisoryText: "reachable", FinishReason: cognitive.FinishCompleted}, nil
}

func TestModelQualificationUsesDiscoveredModelsAndLeavesQualityUnrated(t *testing.T) {
	provider := qualificationProvider{}
	mux := http.NewServeMux()
	mux.Handle("POST /v1/provider-discovery/{provider}/qualify", modelQualificationHandler(map[string]providerRegistration{"gpu": {ProviderInstanceID: "gpu", Provider: provider, Discovery: provider}}, "owner"))
	server := httptest.NewServer(mux)
	defer server.Close()
	request := httptest.NewRequest(http.MethodPost, server.URL+"/v1/provider-discovery/gpu/qualify", strings.NewReader(`{"prompt":"ready","max_output_bytes":128}`))
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()
	mux.ServeHTTP(response, request)
	if response.Code != http.StatusOK {
		t.Fatalf("qualification status=%d body=%s", response.Code, response.Body.String())
	}
	if !strings.Contains(response.Body.String(), `"status":"REACHABLE"`) || !strings.Contains(response.Body.String(), `"category":"UNRATED"`) || strings.Contains(response.Body.String(), `"quality_assessed":true`) {
		t.Fatalf("qualification overclaimed quality: %s", response.Body.String())
	}
}

func TestModelQualificationRejectsUnboundedOutput(t *testing.T) {
	provider := qualificationProvider{}
	mux := http.NewServeMux()
	mux.Handle("POST /v1/provider-discovery/{provider}/qualify", modelQualificationHandler(map[string]providerRegistration{"cpu": {ProviderInstanceID: "cpu", Provider: provider, Discovery: provider}}, "owner"))
	request := httptest.NewRequest(http.MethodPost, "/v1/provider-discovery/cpu/qualify", strings.NewReader(`{"prompt":"ready","max_output_bytes":8192}`))
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()
	mux.ServeHTTP(response, request)
	if response.Code != http.StatusBadRequest {
		t.Fatalf("unbounded qualification status=%d", response.Code)
	}
}

type qualificationObserverStub struct {
	observations []capability.ModelQualificationObservation
}

func (s *qualificationObserverStub) ObserveQualification(_ context.Context, value capability.ModelQualificationObservation) error {
	s.observations = append(s.observations, value)
	return nil
}

func TestModelQualificationPersistsReachabilityWithoutQuality(t *testing.T) {
	provider := qualificationProvider{}
	observer := &qualificationObserverStub{}
	mux := http.NewServeMux()
	mux.Handle("POST /v1/provider-discovery/{provider}/qualify", modelQualificationHandler(map[string]providerRegistration{"cpu": {ProviderInstanceID: "cpu", Provider: provider, Discovery: provider}}, "owner", observer))
	request := httptest.NewRequest(http.MethodPost, "/v1/provider-discovery/cpu/qualify", strings.NewReader(`{"prompt":"ready","max_output_bytes":128,"model_ids":["model-a"]}`))
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()
	mux.ServeHTTP(response, request)
	if response.Code != http.StatusOK || len(observer.observations) != 1 {
		t.Fatalf("qualification persistence status=%d observations=%d body=%s", response.Code, len(observer.observations), response.Body.String())
	}
	if err := observer.observations[0].Validate(); err != nil || observer.observations[0].Status != capability.ModelQualificationReachable {
		t.Fatalf("invalid qualification observation: %+v err=%v", observer.observations[0], err)
	}
}

func TestModelQualificationDeduplicatesSameProviderModelFlight(t *testing.T) {
	coordinator := &qualificationCoordinator{}
	started := make(chan struct{})
	release := make(chan struct{})
	firstDone := make(chan struct{})
	secondDone := make(chan struct{})
	result := modelQualificationResponse{Provider: "medium"}
	go func() {
		_, _, err := coordinator.do(context.Background(), "medium|model-a", func() (modelQualificationResponse, error) {
			close(started)
			<-release
			return result, nil
		})
		if err != nil {
			t.Errorf("leader qualification failed: %v", err)
		}
		close(firstDone)
	}()
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("qualification invocation did not start")
	}
	go func() {
		actual, leader, err := coordinator.do(context.Background(), "medium|model-a", func() (modelQualificationResponse, error) {
			t.Error("duplicate qualification operation ran")
			return modelQualificationResponse{}, nil
		})
		if err != nil || leader || actual.Provider != result.Provider {
			t.Errorf("joined qualification response=%+v leader=%t err=%v", actual, leader, err)
		}
		close(secondDone)
	}()
	deadline := time.Now().Add(time.Second)
	for {
		coordinator.mu.Lock()
		flight := coordinator.flights["medium|model-a"]
		waiters := 0
		if flight != nil {
			waiters = flight.waiters
		}
		coordinator.mu.Unlock()
		if waiters == 1 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("second qualification did not join the in-progress flight")
		}
		time.Sleep(time.Millisecond)
	}
	close(release)
	for _, done := range []chan struct{}{firstDone, secondDone} {
		select {
		case <-done:
		case <-time.After(time.Second):
			t.Fatal("qualification flight did not complete")
		}
	}
}