| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329 |
- //go:build integration
- package repository
- import (
- "context"
- "testing"
- "time"
- dbent "github.com/Wei-Shaw/sub2api/ent"
- "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
- "github.com/Wei-Shaw/sub2api/internal/service"
- "github.com/stretchr/testify/suite"
- )
- type ProxyRepoSuite struct {
- suite.Suite
- ctx context.Context
- tx *dbent.Tx
- repo *proxyRepository
- }
- func (s *ProxyRepoSuite) SetupTest() {
- s.ctx = context.Background()
- tx := testEntTx(s.T())
- s.tx = tx
- s.repo = newProxyRepositoryWithSQL(tx.Client(), tx)
- }
- func TestProxyRepoSuite(t *testing.T) {
- suite.Run(t, new(ProxyRepoSuite))
- }
- // --- Create / GetByID / Update / Delete ---
- func (s *ProxyRepoSuite) TestCreate() {
- proxy := &service.Proxy{
- Name: "test-create",
- Protocol: "http",
- Host: "127.0.0.1",
- Port: 8080,
- Status: service.StatusActive,
- }
- err := s.repo.Create(s.ctx, proxy)
- s.Require().NoError(err, "Create")
- s.Require().NotZero(proxy.ID, "expected ID to be set")
- got, err := s.repo.GetByID(s.ctx, proxy.ID)
- s.Require().NoError(err, "GetByID")
- s.Require().Equal("test-create", got.Name)
- }
- func (s *ProxyRepoSuite) TestGetByID_NotFound() {
- _, err := s.repo.GetByID(s.ctx, 999999)
- s.Require().Error(err, "expected error for non-existent ID")
- }
- func (s *ProxyRepoSuite) TestUpdate() {
- proxy := &service.Proxy{
- Name: "original",
- Protocol: "http",
- Host: "127.0.0.1",
- Port: 8080,
- Status: service.StatusActive,
- }
- s.Require().NoError(s.repo.Create(s.ctx, proxy))
- proxy.Name = "updated"
- err := s.repo.Update(s.ctx, proxy)
- s.Require().NoError(err, "Update")
- got, err := s.repo.GetByID(s.ctx, proxy.ID)
- s.Require().NoError(err, "GetByID after update")
- s.Require().Equal("updated", got.Name)
- }
- func (s *ProxyRepoSuite) TestDelete() {
- proxy := &service.Proxy{
- Name: "to-delete",
- Protocol: "http",
- Host: "127.0.0.1",
- Port: 8080,
- Status: service.StatusActive,
- }
- s.Require().NoError(s.repo.Create(s.ctx, proxy))
- err := s.repo.Delete(s.ctx, proxy.ID)
- s.Require().NoError(err, "Delete")
- _, err = s.repo.GetByID(s.ctx, proxy.ID)
- s.Require().Error(err, "expected error after delete")
- }
- // --- List / ListWithFilters ---
- func (s *ProxyRepoSuite) TestList() {
- s.mustCreateProxy(&service.Proxy{Name: "p1", Protocol: "http", Host: "127.0.0.1", Port: 8080, Status: service.StatusActive})
- s.mustCreateProxy(&service.Proxy{Name: "p2", Protocol: "http", Host: "127.0.0.1", Port: 8081, Status: service.StatusActive})
- proxies, page, err := s.repo.List(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10})
- s.Require().NoError(err, "List")
- s.Require().Len(proxies, 2)
- s.Require().Equal(int64(2), page.Total)
- }
- func (s *ProxyRepoSuite) TestListWithFilters_Protocol() {
- s.mustCreateProxy(&service.Proxy{Name: "p1", Protocol: "http", Host: "127.0.0.1", Port: 8080, Status: service.StatusActive})
- s.mustCreateProxy(&service.Proxy{Name: "p2", Protocol: "socks5", Host: "127.0.0.1", Port: 8081, Status: service.StatusActive})
- proxies, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, "socks5", "", "")
- s.Require().NoError(err)
- s.Require().Len(proxies, 1)
- s.Require().Equal("socks5", proxies[0].Protocol)
- }
- func (s *ProxyRepoSuite) TestListWithFilters_Status() {
- s.mustCreateProxy(&service.Proxy{Name: "p1", Protocol: "http", Host: "127.0.0.1", Port: 8080, Status: service.StatusActive})
- s.mustCreateProxy(&service.Proxy{Name: "p2", Protocol: "http", Host: "127.0.0.1", Port: 8081, Status: service.StatusDisabled})
- proxies, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, "", service.StatusDisabled, "")
- s.Require().NoError(err)
- s.Require().Len(proxies, 1)
- s.Require().Equal(service.StatusDisabled, proxies[0].Status)
- }
- func (s *ProxyRepoSuite) TestListWithFilters_Search() {
- s.mustCreateProxy(&service.Proxy{Name: "production-proxy", Protocol: "http", Host: "127.0.0.1", Port: 8080, Status: service.StatusActive})
- s.mustCreateProxy(&service.Proxy{Name: "dev-proxy", Protocol: "http", Host: "127.0.0.1", Port: 8081, Status: service.StatusActive})
- proxies, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{Page: 1, PageSize: 10}, "", "", "prod")
- s.Require().NoError(err)
- s.Require().Len(proxies, 1)
- s.Require().Contains(proxies[0].Name, "production")
- }
- // --- ListActive ---
- func (s *ProxyRepoSuite) TestListActive() {
- s.mustCreateProxy(&service.Proxy{Name: "active1", Protocol: "http", Host: "127.0.0.1", Port: 8080, Status: service.StatusActive})
- s.mustCreateProxy(&service.Proxy{Name: "inactive1", Protocol: "http", Host: "127.0.0.1", Port: 8081, Status: service.StatusDisabled})
- proxies, err := s.repo.ListActive(s.ctx)
- s.Require().NoError(err, "ListActive")
- s.Require().Len(proxies, 1)
- s.Require().Equal("active1", proxies[0].Name)
- }
- // --- ExistsByHostPortAuth ---
- func (s *ProxyRepoSuite) TestExistsByHostPortAuth() {
- s.mustCreateProxy(&service.Proxy{
- Name: "p1",
- Protocol: "http",
- Host: "1.2.3.4",
- Port: 8080,
- Username: "user",
- Password: "pass",
- Status: service.StatusActive,
- })
- exists, err := s.repo.ExistsByHostPortAuth(s.ctx, "1.2.3.4", 8080, "user", "pass")
- s.Require().NoError(err, "ExistsByHostPortAuth")
- s.Require().True(exists)
- notExists, err := s.repo.ExistsByHostPortAuth(s.ctx, "1.2.3.4", 8080, "wrong", "creds")
- s.Require().NoError(err)
- s.Require().False(notExists)
- }
- func (s *ProxyRepoSuite) TestExistsByHostPortAuth_NoAuth() {
- s.mustCreateProxy(&service.Proxy{
- Name: "p-noauth",
- Protocol: "http",
- Host: "5.6.7.8",
- Port: 8081,
- Username: "",
- Password: "",
- Status: service.StatusActive,
- })
- exists, err := s.repo.ExistsByHostPortAuth(s.ctx, "5.6.7.8", 8081, "", "")
- s.Require().NoError(err)
- s.Require().True(exists)
- }
- // --- CountAccountsByProxyID ---
- func (s *ProxyRepoSuite) TestCountAccountsByProxyID() {
- proxy := s.mustCreateProxy(&service.Proxy{Name: "p-count", Protocol: "http", Host: "127.0.0.1", Port: 8080, Status: service.StatusActive})
- s.mustInsertAccount("a1", &proxy.ID)
- s.mustInsertAccount("a2", &proxy.ID)
- s.mustInsertAccount("a3", nil) // no proxy
- count, err := s.repo.CountAccountsByProxyID(s.ctx, proxy.ID)
- s.Require().NoError(err, "CountAccountsByProxyID")
- s.Require().Equal(int64(2), count)
- }
- func (s *ProxyRepoSuite) TestCountAccountsByProxyID_Zero() {
- proxy := s.mustCreateProxy(&service.Proxy{Name: "p-zero", Protocol: "http", Host: "127.0.0.1", Port: 8080, Status: service.StatusActive})
- count, err := s.repo.CountAccountsByProxyID(s.ctx, proxy.ID)
- s.Require().NoError(err)
- s.Require().Zero(count)
- }
- // --- GetAccountCountsForProxies ---
- func (s *ProxyRepoSuite) TestGetAccountCountsForProxies() {
- p1 := s.mustCreateProxy(&service.Proxy{Name: "p1", Protocol: "http", Host: "127.0.0.1", Port: 8080, Status: service.StatusActive})
- p2 := s.mustCreateProxy(&service.Proxy{Name: "p2", Protocol: "http", Host: "127.0.0.1", Port: 8081, Status: service.StatusActive})
- s.mustInsertAccount("a1", &p1.ID)
- s.mustInsertAccount("a2", &p1.ID)
- s.mustInsertAccount("a3", &p2.ID)
- counts, err := s.repo.GetAccountCountsForProxies(s.ctx)
- s.Require().NoError(err, "GetAccountCountsForProxies")
- s.Require().Equal(int64(2), counts[p1.ID])
- s.Require().Equal(int64(1), counts[p2.ID])
- }
- func (s *ProxyRepoSuite) TestGetAccountCountsForProxies_Empty() {
- counts, err := s.repo.GetAccountCountsForProxies(s.ctx)
- s.Require().NoError(err)
- s.Require().Empty(counts)
- }
- // --- ListActiveWithAccountCount ---
- func (s *ProxyRepoSuite) TestListActiveWithAccountCount() {
- base := time.Date(2025, 1, 1, 12, 0, 0, 0, time.UTC)
- p1 := s.mustCreateProxyWithTimes("p1", service.StatusActive, base.Add(-1*time.Hour))
- p2 := s.mustCreateProxyWithTimes("p2", service.StatusActive, base)
- s.mustCreateProxyWithTimes("p3-inactive", service.StatusDisabled, base.Add(1*time.Hour))
- s.mustInsertAccount("a1", &p1.ID)
- s.mustInsertAccount("a2", &p1.ID)
- s.mustInsertAccount("a3", &p2.ID)
- withCounts, err := s.repo.ListActiveWithAccountCount(s.ctx)
- s.Require().NoError(err, "ListActiveWithAccountCount")
- s.Require().Len(withCounts, 2, "expected 2 active proxies")
- // Sorted by created_at DESC, so p2 first
- s.Require().Equal(p2.ID, withCounts[0].ID)
- s.Require().Equal(int64(1), withCounts[0].AccountCount)
- s.Require().Equal(p1.ID, withCounts[1].ID)
- s.Require().Equal(int64(2), withCounts[1].AccountCount)
- }
- // --- Combined original test ---
- func (s *ProxyRepoSuite) TestExistsByHostPortAuth_And_AccountCountAggregates() {
- p1 := s.mustCreateProxy(&service.Proxy{Name: "p1", Protocol: "http", Host: "1.2.3.4", Port: 8080, Username: "u", Password: "p", Status: service.StatusActive})
- p2 := s.mustCreateProxy(&service.Proxy{Name: "p2", Protocol: "http", Host: "5.6.7.8", Port: 8081, Username: "", Password: "", Status: service.StatusActive})
- exists, err := s.repo.ExistsByHostPortAuth(s.ctx, "1.2.3.4", 8080, "u", "p")
- s.Require().NoError(err, "ExistsByHostPortAuth")
- s.Require().True(exists, "expected proxy to exist")
- s.mustInsertAccount("a1", &p1.ID)
- s.mustInsertAccount("a2", &p1.ID)
- s.mustInsertAccount("a3", &p2.ID)
- count1, err := s.repo.CountAccountsByProxyID(s.ctx, p1.ID)
- s.Require().NoError(err, "CountAccountsByProxyID")
- s.Require().Equal(int64(2), count1, "expected 2 accounts for p1")
- counts, err := s.repo.GetAccountCountsForProxies(s.ctx)
- s.Require().NoError(err, "GetAccountCountsForProxies")
- s.Require().Equal(int64(2), counts[p1.ID])
- s.Require().Equal(int64(1), counts[p2.ID])
- withCounts, err := s.repo.ListActiveWithAccountCount(s.ctx)
- s.Require().NoError(err, "ListActiveWithAccountCount")
- s.Require().Len(withCounts, 2, "expected 2 proxies")
- for _, pc := range withCounts {
- switch pc.ID {
- case p1.ID:
- s.Require().Equal(int64(2), pc.AccountCount, "p1 count mismatch")
- case p2.ID:
- s.Require().Equal(int64(1), pc.AccountCount, "p2 count mismatch")
- default:
- s.Require().Fail("unexpected proxy id", pc.ID)
- }
- }
- }
- func (s *ProxyRepoSuite) mustCreateProxy(p *service.Proxy) *service.Proxy {
- s.T().Helper()
- s.Require().NoError(s.repo.Create(s.ctx, p), "create proxy")
- return p
- }
- func (s *ProxyRepoSuite) mustCreateProxyWithTimes(name, status string, createdAt time.Time) *service.Proxy {
- s.T().Helper()
- // Use the repository create for standard fields, then update timestamps via raw SQL to keep deterministic ordering.
- p := s.mustCreateProxy(&service.Proxy{
- Name: name,
- Protocol: "http",
- Host: "127.0.0.1",
- Port: 8080,
- Status: status,
- })
- _, err := s.tx.ExecContext(s.ctx, "UPDATE proxies SET created_at = $1, updated_at = $1 WHERE id = $2", createdAt, p.ID)
- s.Require().NoError(err, "update proxy timestamps")
- return p
- }
- func (s *ProxyRepoSuite) mustInsertAccount(name string, proxyID *int64) {
- s.T().Helper()
- var pid any
- if proxyID != nil {
- pid = *proxyID
- }
- _, err := s.tx.ExecContext(
- s.ctx,
- "INSERT INTO accounts (name, platform, type, proxy_id) VALUES ($1, $2, $3, $4)",
- name,
- service.PlatformAnthropic,
- service.AccountTypeOAuth,
- pid,
- )
- s.Require().NoError(err, "insert account")
- }
|