heicode-mananger/heicode/controller/agent_template_test.go

package controller

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/heicode/manager/model"
	"github.com/stretchr/testify/require"
)

func TestBuildAgentEnvFromBindings_BuiltinBlobMetadataOnly(t *testing.T) {
	setupResourceControllerTestDB(t)

	// No secret_ref -> secret fields (BLOB_KEY) are skipped, not an error.
	b := model.ResourceBinding{
		UserId:       7,
		Name:         "blob1",
		ResourceType: "blob",
		Provider:     "blob",
		Metadata:     `{"account":"myacct","container":"uploads"}`,
	}
	require.NoError(t, model.DB.Create(&b).Error)

	env, err := buildAgentEnvFromBindings(7, []int{b.Id})
	require.NoError(t, err)
	require.Equal(t, "myacct", env["BLOB_ACCOUNT"])
	require.Equal(t, "uploads", env["BLOB_CONTAINER"])
	_, hasKey := env["BLOB_KEY"]
	require.False(t, hasKey) // no secret bound
}

func TestBuildAgentEnvFromBindings_DatabaseProviderPrefix(t *testing.T) {
	setupResourceControllerTestDB(t)
	cases := []struct {
		provider string
		wantHost string // env var that should carry the host
	}{
		{"mysql", "MYSQL_HOST"},
		{"pg", "PG_HOST"},
		{"postgresql", "PG_HOST"}, // alias normalises to pg
		{"redis", "REDIS_HOST"},
		{"mg", "MONGO_HOST"}, // alias normalises to mongo
	}
	for _, tc := range cases {
		b := model.ResourceBinding{
			UserId:       7,
			Name:         "db-" + tc.provider,
			ResourceType: "database",
			Provider:     tc.provider,
			Metadata:     `{"host":"db.example","port":"5432","database":"app","user":"u"}`,
		}
		require.NoError(t, model.DB.Create(&b).Error)
		env, err := buildAgentEnvFromBindings(7, []int{b.Id})
		require.NoError(t, err, tc.provider)
		require.Equal(t, "db.example", env[tc.wantHost], "provider %s -> %s", tc.provider, tc.wantHost)
	}
}

func TestBuildAgentEnvFromBindings_GitNamesProviderAgnostic(t *testing.T) {
	setupResourceControllerTestDB(t)
	for _, provider := range []string{"github", "gitea", "gitlab"} {
		b := model.ResourceBinding{
			UserId:       7,
			Name:         "git-" + provider,
			ResourceType: "git",
			Provider:     provider,
			Metadata:     `{"provider":"` + provider + `","repo_url":"https://x/owner/repo","default_branch":"main"}`,
		}
		require.NoError(t, model.DB.Create(&b).Error)
		env, err := buildAgentEnvFromBindings(7, []int{b.Id})
		require.NoError(t, err)
		require.Equal(t, "https://x/owner/repo", env["GIT_REPO_URL"])
		require.Equal(t, "main", env["GIT_BRANCH"])
		require.Equal(t, provider, env["GIT_PROVIDER"])
	}
}

func TestBuildAgentEnvFromBindings_OwnershipEnforced(t *testing.T) {
	setupResourceControllerTestDB(t)
	b := model.ResourceBinding{UserId: 7, Name: "x", ResourceType: "blob", Provider: "blob", Metadata: `{}`}
	require.NoError(t, model.DB.Create(&b).Error)

	_, err := buildAgentEnvFromBindings(99, []int{b.Id}) // different user
	require.Error(t, err)
}

func TestBuildAgentEnvFromBindings_UnsupportedType(t *testing.T) {
	setupResourceControllerTestDB(t)
	b := model.ResourceBinding{UserId: 7, Name: "x", ResourceType: "quantum", Provider: "weird", Metadata: `{}`}
	require.NoError(t, model.DB.Create(&b).Error)

	_, err := buildAgentEnvFromBindings(7, []int{b.Id})
	require.Error(t, err)
}

func TestBuildAgentEnvFromBindings_Empty(t *testing.T) {
	setupResourceControllerTestDB(t)
	env, err := buildAgentEnvFromBindings(7, nil)
	require.NoError(t, err)
	require.Empty(t, env)
}

func TestAgentTemplatePathSubstitution(t *testing.T) {
	require.Contains(t, agentTemplateStartPath("tpl1"), "/templates/tpl1/start")
	require.Contains(t, agentTemplateAgentStopPath("rt-9"), "/agents/rt-9/stop")
	require.Contains(t, agentTemplateAgentPath("rt-9"), "/agents/rt-9")
}

func TestAMStartTemplateAgent_RoundTrip(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, http.MethodPost, r.Method)
		require.Equal(t, "/api/agent/templates/tpl1/start", r.URL.Path)
		body, _ := io.ReadAll(r.Body)
		require.Contains(t, string(body), "manager_deployment_id")
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"success":true,"data":{"runtime_id":"rt-9","subdomain":"https://x.example","access_token":"tok","status":"running"}}`))
	}))
	defer srv.Close()
	t.Setenv("AGENT_RUNTIME_BASE_URL", srv.URL)

	res, err := amStartTemplateAgent(context.Background(), "tpl1", "dep_1", map[string]string{"K": "V"}, "https://cb")
	require.NoError(t, err)
	require.Equal(t, "rt-9", res.RuntimeID)
	require.Equal(t, "https://x.example", res.Subdomain)
	require.Equal(t, "tok", res.AccessToken)
	require.Equal(t, "running", res.Status)
}

func TestAMGetAgentStatus_RoundTrip(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.Equal(t, "/api/agent/agents/rt-1", r.URL.Path)
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"success":true,"data":{"status":"running"}}`))
	}))
	defer srv.Close()
	t.Setenv("AGENT_RUNTIME_BASE_URL", srv.URL)

	status, err := amGetAgentStatus(context.Background(), "rt-1")
	require.NoError(t, err)
	require.Equal(t, "running", status)
}

func TestTemplateAgentResponse(t *testing.T) {
	row := model.AgentDeployment{
		DeploymentID:   "dep_abc",
		TemplateID:     "tpl1",
		Subdomain:      "https://abc.agents.example",
		AccessToken:    "tok",
		BindingIDsJSON: "[1,2]",
		Status:         "running",
	}
	resp := templateAgentResponse(row)
	require.Equal(t, "dep_abc", resp["agent_id"])
	require.Equal(t, "tpl1", resp["template_id"])
	require.Equal(t, []int{1, 2}, resp["binding_ids"])
	require.Equal(t, "running", resp["status"])
}