mirror of
https://github.com/google/ax.git
synced 2026-10-31 07:33:11 +08:00
JBD
dc4f36cdba
AX is undergoing a significant redesign, moving away from a single CLI with an embedded Python harness toward a general-purpose orchestration layer for agentic tasks. The repository is reorganized around three binaries: ax-server exposes a gRPC API for submitting Task manifests, ax-controller runs as a horizontally scaled reconciler that consumes work from Redis Streams, and ax-task-runner executes tasks inside sandboxed workers. Task state now lives in Redis rather than Kubernetes CRDs so the system can handle millions of short-lived tasks without straining etcd. The legacy Python harness, ATE client, SQL event log, and skill examples are removed, and new design, concepts, manifest, networking, runner, and sandbox docs plus Kubernetes deploy manifests describe the new architecture.
287 lines
8.6 KiB
Go
287 lines
8.6 KiB
Go
// Copyright 2026 Google LLC
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package workspace_test
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/ax/internal/workspace"
|
|
"github.com/google/ax/pkg/apis/v1alpha1"
|
|
)
|
|
|
|
func TestSetupWorkspace_MaidenRun(t *testing.T) {
|
|
tempDir, err := os.MkdirTemp("", "ax-ws-test-*")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(tempDir)
|
|
|
|
ws := &v1alpha1.Workspace{
|
|
Metadata: &v1alpha1.ObjectMeta{
|
|
Name: "test-ws",
|
|
},
|
|
Spec: &v1alpha1.WorkspaceSpec{
|
|
Skills: &v1alpha1.SkillsConfig{
|
|
Path: filepath.Join(tempDir, "skills"),
|
|
},
|
|
},
|
|
}
|
|
|
|
stateDir := filepath.Join(tempDir, "ax-state")
|
|
origAXDir := workspace.AXDir
|
|
workspace.AXDir = stateDir
|
|
defer func() { workspace.AXDir = origAXDir }()
|
|
|
|
// 1. Maiden run
|
|
res1, err := workspace.SetupWorkspace(context.Background(), ws, tempDir, "Test maiden run setup")
|
|
if err != nil {
|
|
t.Fatalf("maiden setup failed: %v", err)
|
|
}
|
|
if !res1.IsMaidenRun {
|
|
t.Errorf("expected maiden run to be true")
|
|
}
|
|
|
|
// Verify marker file exists under stateDir/initialized
|
|
markerFile := filepath.Join(stateDir, workspace.MarkerName(tempDir))
|
|
if _, err := os.Stat(markerFile); os.IsNotExist(err) {
|
|
t.Errorf("expected marker file %s to exist", markerFile)
|
|
}
|
|
|
|
// Verify workspace directory has no .ax-initialized or .ax directory
|
|
if _, err := os.Stat(filepath.Join(tempDir, ".ax-initialized")); !os.IsNotExist(err) {
|
|
t.Errorf("expected .ax-initialized not to be in workspace")
|
|
}
|
|
if _, err := os.Stat(filepath.Join(tempDir, ".ax")); !os.IsNotExist(err) {
|
|
t.Errorf("expected .ax directory not to be in workspace")
|
|
}
|
|
|
|
// 2. Second run (should detect already initialized)
|
|
res2, err := workspace.SetupWorkspace(context.Background(), ws, tempDir, "")
|
|
if err != nil {
|
|
t.Fatalf("second setup failed: %v", err)
|
|
}
|
|
if res2.IsMaidenRun {
|
|
t.Errorf("expected maiden run to be false on second run")
|
|
}
|
|
}
|
|
|
|
func TestSetupWorkspace_GitSubdir(t *testing.T) {
|
|
tempDir, err := os.MkdirTemp("", "ax-ws-git-*")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(tempDir)
|
|
|
|
// Create a dummy source git repository
|
|
srcDir, err := os.MkdirTemp("", "ax-src-git-*")
|
|
if err != nil {
|
|
t.Fatalf("failed to create src dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(srcDir)
|
|
|
|
// Init source repo
|
|
// git init -b requires git 2.28+; symbolic-ref works on any version.
|
|
initCmd := exec.Command("sh", "-c", "git init && git symbolic-ref HEAD refs/heads/main && git config user.email 'test@ax.io' && git config user.name 'AX' && echo 'hello' > README.md && git add . && git commit -m 'initial'")
|
|
initCmd.Dir = srcDir
|
|
if out, err := initCmd.CombinedOutput(); err != nil {
|
|
t.Fatalf("failed to init test git repo: %v (%s)", err, string(out))
|
|
}
|
|
|
|
ws := &v1alpha1.Workspace{
|
|
Metadata: &v1alpha1.ObjectMeta{
|
|
Name: "test-git-ws",
|
|
},
|
|
Spec: &v1alpha1.WorkspaceSpec{
|
|
Git: []*v1alpha1.GitRepo{
|
|
{
|
|
Name: "ax",
|
|
Repo: srcDir,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
stateDir := filepath.Join(tempDir, "ax-state")
|
|
origAXDir := workspace.AXDir
|
|
workspace.AXDir = stateDir
|
|
defer func() { workspace.AXDir = origAXDir }()
|
|
|
|
res, err := workspace.SetupWorkspace(context.Background(), ws, tempDir, "")
|
|
if err != nil {
|
|
t.Fatalf("setup failed: %v", err)
|
|
}
|
|
if len(res.ClonedRepos) != 1 {
|
|
t.Errorf("expected 1 cloned repo, got %d", len(res.ClonedRepos))
|
|
}
|
|
|
|
// Verify repo was cloned into subdirectory 'ax' inside tempDir
|
|
clonedGit := filepath.Join(tempDir, "ax", ".git")
|
|
if _, err := os.Stat(clonedGit); os.IsNotExist(err) {
|
|
t.Errorf("expected %s to exist", clonedGit)
|
|
}
|
|
|
|
// Verify git-success.log is created under stateDir and NOT inside workspace
|
|
logFile := filepath.Join(stateDir, "git-success.log")
|
|
if _, err := os.Stat(logFile); os.IsNotExist(err) {
|
|
t.Errorf("expected git log file %s to exist in stateDir", logFile)
|
|
}
|
|
if _, err := os.Stat(filepath.Join(tempDir, ".ax-git-success.log")); !os.IsNotExist(err) {
|
|
t.Errorf("expected .ax-git-success.log NOT to exist in workspace")
|
|
}
|
|
}
|
|
|
|
func TestSetupWorkspace_GitError(t *testing.T) {
|
|
tempDir, err := os.MkdirTemp("", "ax-ws-git-err-*")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(tempDir)
|
|
|
|
ws := &v1alpha1.Workspace{
|
|
Metadata: &v1alpha1.ObjectMeta{
|
|
Name: "test-git-err-ws",
|
|
},
|
|
Spec: &v1alpha1.WorkspaceSpec{
|
|
Git: []*v1alpha1.GitRepo{
|
|
{
|
|
Name: "invalid-repo",
|
|
Repo: "https://127.0.0.1:9/nonexistent/repo.git",
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
stateDir := filepath.Join(tempDir, "ax-state")
|
|
origAXDir := workspace.AXDir
|
|
workspace.AXDir = stateDir
|
|
defer func() { workspace.AXDir = origAXDir }()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
|
defer cancel()
|
|
|
|
_, _ = workspace.SetupWorkspace(ctx, ws, tempDir, "")
|
|
|
|
// Verify git-error.log is created under stateDir and NOT inside workspace
|
|
errLog := filepath.Join(stateDir, "git-error.log")
|
|
if _, err := os.Stat(errLog); os.IsNotExist(err) {
|
|
t.Errorf("expected git error log %s to exist under stateDir", errLog)
|
|
}
|
|
|
|
// Verify workspace directory does not have .ax-git-error.log or .ax
|
|
if _, err := os.Stat(filepath.Join(tempDir, ".ax-git-error.log")); !os.IsNotExist(err) {
|
|
t.Errorf("expected .ax-git-error.log NOT to exist in workspace")
|
|
}
|
|
if _, err := os.Stat(filepath.Join(tempDir, ".ax")); !os.IsNotExist(err) {
|
|
t.Errorf("expected .ax directory NOT to exist in workspace")
|
|
}
|
|
|
|
// Verify initialization marker was not created due to failure
|
|
markerFile := filepath.Join(stateDir, workspace.MarkerName(tempDir))
|
|
if _, err := os.Stat(markerFile); !os.IsNotExist(err) {
|
|
t.Errorf("expected marker file %s NOT to exist after failure", markerFile)
|
|
}
|
|
}
|
|
|
|
func TestSetupWorkspace_GitDepth(t *testing.T) {
|
|
tempDir, err := os.MkdirTemp("", "ax-ws-git-depth-*")
|
|
if err != nil {
|
|
t.Fatalf("failed to create temp dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(tempDir)
|
|
|
|
// Create a dummy source git repository with 3 commits
|
|
srcDir, err := os.MkdirTemp("", "ax-src-git-*")
|
|
if err != nil {
|
|
t.Fatalf("failed to create src dir: %v", err)
|
|
}
|
|
defer os.RemoveAll(srcDir)
|
|
|
|
initCmd := exec.Command("sh", "-c", `
|
|
git init &&
|
|
git symbolic-ref HEAD refs/heads/main &&
|
|
git config user.email 'test@ax.io' &&
|
|
git config user.name 'AX' &&
|
|
echo '1' > file.txt && git add . && git commit -m 'commit 1' &&
|
|
echo '2' > file.txt && git add . && git commit -m 'commit 2' &&
|
|
echo '3' > file.txt && git add . && git commit -m 'commit 3'
|
|
`)
|
|
initCmd.Dir = srcDir
|
|
if out, err := initCmd.CombinedOutput(); err != nil {
|
|
t.Fatalf("failed to init test git repo with commits: %v (%s)", err, string(out))
|
|
}
|
|
|
|
ws := &v1alpha1.Workspace{
|
|
Metadata: &v1alpha1.ObjectMeta{
|
|
Name: "test-shallow-git-ws",
|
|
},
|
|
Spec: &v1alpha1.WorkspaceSpec{
|
|
Git: []*v1alpha1.GitRepo{
|
|
{
|
|
Name: "ax",
|
|
Repo: srcDir,
|
|
Depth: 1,
|
|
},
|
|
},
|
|
},
|
|
}
|
|
|
|
stateDir := filepath.Join(tempDir, "ax-state")
|
|
origAXDir := workspace.AXDir
|
|
workspace.AXDir = stateDir
|
|
defer func() { workspace.AXDir = origAXDir }()
|
|
|
|
res, err := workspace.SetupWorkspace(context.Background(), ws, tempDir, "")
|
|
if err != nil {
|
|
t.Fatalf("setup failed: %v", err)
|
|
}
|
|
if len(res.ClonedRepos) != 1 {
|
|
t.Errorf("expected 1 cloned repo, got %d", len(res.ClonedRepos))
|
|
}
|
|
|
|
// Verify commit count in cloned repo is exactly 1
|
|
countCmd := exec.Command("git", "rev-list", "--count", "HEAD")
|
|
countCmd.Dir = filepath.Join(tempDir, "ax")
|
|
out, err := countCmd.Output()
|
|
if err != nil {
|
|
t.Fatalf("failed to count commits: %v", err)
|
|
}
|
|
if got := string(out); got != "1\n" && got != "1" {
|
|
t.Errorf("expected 1 commit with depth=1, got %q", got)
|
|
}
|
|
}
|
|
|
|
func TestMarkerName(t *testing.T) {
|
|
tests := map[string]string{
|
|
"/workspace": "initialized-workspace",
|
|
"/workspace/": "initialized-workspace",
|
|
"": "initialized-workspace",
|
|
"/": "initialized-root",
|
|
"/workspace/tools": "initialized-workspace-tools",
|
|
"/workspace/a/b": "initialized-workspace-a-b",
|
|
"/srv/data": "initialized-srv-data",
|
|
"/workspace/tools/": "initialized-workspace-tools",
|
|
}
|
|
for path, want := range tests {
|
|
if got := workspace.MarkerName(path); got != want {
|
|
t.Errorf("MarkerName(%q) = %q, want %q", path, got, want)
|
|
}
|
|
}
|
|
}
|