package controllers

import (
	"testing"

	"apiserver/models"
	"apiserver/store"
)

func TestCountRunningJobsPerHost(t *testing.T) {
	store.InitStores()

	// RUN job on t1.
	store.Jobs.Put(1, models.JobInfo{
		JobStatus: models.JOBSTAT["RUN"],
		ExecHosts: models.Jsonb{String: `{"host_slots":[{"host_name":"t1","num_slots":1}],"num_hosts":1}`},
	})
	// RUN|USUSP job on t2 — the RUN bit must still count it.
	store.Jobs.Put(2, models.JobInfo{
		JobStatus: models.JOBSTAT["RUN"] | models.JOBSTAT["USUSP"],
		ExecHosts: models.Jsonb{String: `{"host_slots":[{"host_name":"t2","num_slots":1}],"num_hosts":1}`},
	})
	// Multi-host RUN job: one job on t1 + t3 counts on both.
	store.Jobs.Put(3, models.JobInfo{
		JobStatus: models.JOBSTAT["RUN"],
		ExecHosts: models.Jsonb{String: `{"host_slots":[{"host_name":"t1","num_slots":1},{"host_name":"t3","num_slots":1}],"num_hosts":2}`},
	})
	// PEND and DONE jobs must NOT count.
	store.Jobs.Put(4, models.JobInfo{
		JobStatus: models.JOBSTAT["PEND"],
		ExecHosts: models.Jsonb{String: `{"host_slots":[{"host_name":"t3","num_slots":1}],"num_hosts":1}`},
	})
	store.Jobs.Put(5, models.JobInfo{
		JobStatus: models.JOBSTAT["DONE"],
		ExecHosts: models.Jsonb{String: `{"host_slots":[{"host_name":"t1","num_slots":1}],"num_hosts":1}`},
	})
	// RUN job with empty/bad exec_hosts must not crash or count.
	store.Jobs.Put(6, models.JobInfo{
		JobStatus: models.JOBSTAT["RUN"],
		ExecHosts: models.Jsonb{String: ``},
	})
	store.Jobs.Put(7, models.JobInfo{
		JobStatus: models.JOBSTAT["RUN"],
		ExecHosts: models.Jsonb{String: `not-json`},
	})

	counts := countRunningJobsPerHost()

	if counts["t1"] != 2 { // jobs 1 + 3
		t.Errorf("t1 running = %d, want 2", counts["t1"])
	}
	if counts["t2"] != 1 {
		t.Errorf("t2 running = %d, want 1", counts["t2"])
	}
	if counts["t3"] != 1 { // job 3 only (job 4 is PEND)
		t.Errorf("t3 running = %d, want 1", counts["t3"])
	}
	if _, ok := counts["nohost"]; ok {
		t.Error("unexpected host 'nohost' in counts")
	}
}