<commit_before>package gormigrate

import (
	"errors"
	"fmt"

	"github.com/jinzhu/gorm"
)

const (
	initSchemaMigrationId = "SCHEMA_INIT"
)

// MigrateFunc is the func signature for migrating.
type MigrateFunc func(*gorm.DB) error

// RollbackFunc is the func signature for rollbacking.
type RollbackFunc func(*gorm.DB) error

// InitSchemaFunc is the func signature for initializing the schema.
type InitSchemaFunc func(*gorm.DB) error

// Options define options for all migrations.
type Options struct {
	// TableName is the migration table.
	TableName string
	// IDColumnName is the name of column where the migration id will be stored.
	IDColumnName string
	// IDColumnSize is the length of the migration id column
	IDColumnSize int
	// UseTransaction makes Gormigrate execute migrations inside a single transaction.
	// Keep in mind that not all databases support DDL commands inside transactions.
	UseTransaction bool
}

// Migration represents a database migration (a modification to be made on the database).
type Migration struct {
	// ID is the migration identifier. Usually a timestamp like "201601021504".
	ID string
	// Migrate is a function that will br executed while running this migration.
	Migrate MigrateFunc
	// Rollback will be executed on rollback. Can be nil.
	Rollback RollbackFunc
}

// Gormigrate represents a collection of all migrations of a database schema.
type Gormigrate struct {
	db         *gorm.DB
	tx         *gorm.DB
	options    *Options
	migrations []*Migration
	initSchema InitSchemaFunc
}

// ReservedIDError is returned when a migration is using a reserved ID
type ReservedIDError struct {
	ID string
}

func (e *ReservedIDError) Error() string {
	return fmt.Sprintf(`gormigrate: Reserved migration ID: "%s"`, e.ID)
}

// DuplicatedIDError is returned when more than one migration have the same ID
type DuplicatedIDError struct {
	ID string
}

func (e *DuplicatedIDError) Error() string {
	return fmt.Sprintf(`gormigrate: Duplicated migration ID: "%s"`, e.ID)
}

var (
	// DefaultOptions can be used if you don't want to think about options.
	DefaultOptions = &Options{
		TableName:      "migrations",
		IDColumnName:   "id",
		IDColumnSize:   255,
		UseTransaction: false,
	}

	// ErrRollbackImpossible is returned when trying to rollback a migration
	// that has no rollback function.
	ErrRollbackImpossible = errors.New("gormigrate: It's impossible to rollback this migration")

	// ErrNoMigrationDefined is returned when no migration is defined.
	ErrNoMigrationDefined = errors.New("gormigrate: No migration defined")

	// ErrMissingID is returned when the ID od migration is equal to ""
	ErrMissingID = errors.New("gormigrate: Missing ID in migration")

	// ErrNoRunMigration is returned when any run migration was found while
	// running RollbackLast
	ErrNoRunMigration = errors.New("gormigrate: Could not find last run migration")

	// ErrMigrationIDDoesNotExist is returned when migrating or rolling back to a migration ID that
	// does not exist in the list of migrations
	ErrMigrationIDDoesNotExist = errors.New("gormigrate: Tried to migrate to an ID that doesn't exist")
)

// New returns a new Gormigrate.
func New(db *gorm.DB, options *Options, migrations []*Migration) *Gormigrate {
	if options.TableName == "" {
		options.TableName = DefaultOptions.TableName
	}
	if options.IDColumnName == "" {
		options.IDColumnName = DefaultOptions.IDColumnName
	}
	if options.IDColumnSize == 0 {
		options.IDColumnSize = DefaultOptions.IDColumnSize
	}
	return &Gormigrate{
		db:         db,
		options:    options,
		migrations: migrations,
	}
}

// InitSchema sets a function that is run if no migration is found.
// The idea is preventing to run all migrations when a new clean database
// is being migrating. In this function you should create all tables and
// foreign key necessary to your application.
func (g *Gormigrate) InitSchema(initSchema InitSchemaFunc) {
	g.initSchema = initSchema
}

// Migrate executes all migrations that did not run yet.
func (g *Gormigrate) Migrate() error {
	if !g.hasMigrations() {
		return ErrNoMigrationDefined
	}
	var targetMigrationID string
	if len(g.migrations) > 0 {
		targetMigrationID = g.migrations[len(g.migrations)-1].ID
	}
	return g.migrate(targetMigrationID)
}

// MigrateTo executes all migrations that did not run yet up to the migration that matches `migrationID`.
func (g *Gormigrate) MigrateTo(migrationID string) error {
	if err := g.checkIDExist(migrationID); err != nil {
		return err
	}
	return g.migrate(migrationID)
}

func (g *Gormigrate) migrate(migrationID string) error {
	if !g.hasMigrations() {
		return ErrNoMigrationDefined
	}

	if err := g.checkReservedID(); err != nil {
		return err
	}

	if err := g.checkDuplicatedID(); err != nil {
		return err
	}

	g.begin()
	defer g.rollback()

	if err := g.createMigrationTableIfNotExists(); err != nil {
		return err
	}

	if g.initSchema != nil && g.canInitializeSchema() {
		if err := g.runInitSchema(); err != nil {
			return err
		}
		return g.commit()
	}

	for _, migration := range g.migrations {
		if err := g.runMigration(migration); err != nil {
			return err
		}
		if migrationID != "" && migration.ID == migrationID {
			break
		}
	}

	return g.commit()
}

// There are migrations to apply if either there's a defined
// initSchema function or if the list of migrations is not empty.
func (g *Gormigrate) hasMigrations() bool {
	return g.initSchema != nil || len(g.migrations) > 0
}

// Check whether any migration is using a reserved ID.
// For now there's only have one reserved ID, but there may be more in the future.
func (g *Gormigrate) checkReservedID() error {
	for _, m := range g.migrations {
		if m.ID == initSchemaMigrationId {
			return &ReservedIDError{ID: m.ID}
		}
	}
	return nil
}

func (g *Gormigrate) checkDuplicatedID() error {
	lookup := make(map[string]struct{}, len(g.migrations))
	for _, m := range g.migrations {
		if _, ok := lookup[m.ID]; ok {
			return &DuplicatedIDError{ID: m.ID}
		}
		lookup[m.ID] = struct{}{}
	}
	return nil
}

func (g *Gormigrate) checkIDExist(migrationID string) error {
	for _, migrate := range g.migrations {
		if migrate.ID == migrationID {
			return nil
		}
	}
	return ErrMigrationIDDoesNotExist
}

// RollbackLast undo the last migration
func (g *Gormigrate) RollbackLast() error {
	if len(g.migrations) == 0 {
		return ErrNoMigrationDefined
	}

	g.begin()
	defer g.rollback()

	lastRunMigration, err := g.getLastRunMigration()
	if err != nil {
		return err
	}

	if err := g.rollbackMigration(lastRunMigration); err != nil {
		return err
	}
	return g.commit()
}

// RollbackTo undoes migrations up to the given migration that matches the `migrationID`.
// Migration with the matching `migrationID` is not rolled back.
func (g *Gormigrate) RollbackTo(migrationID string) error {
	if len(g.migrations) == 0 {
		return ErrNoMigrationDefined
	}

	if err := g.checkIDExist(migrationID); err != nil {
		return err
	}

	g.begin()
	defer g.rollback()

	for i := len(g.migrations) - 1; i >= 0; i-- {
		migration := g.migrations[i]
		if migration.ID == migrationID {
			break
		}
		if g.migrationDidRun(migration) {
			if err := g.rollbackMigration(migration); err != nil {
				return err
			}
		}
	}
	return g.commit()
}

func (g *Gormigrate) getLastRunMigration() (*Migration, error) {
	for i := len(g.migrations) - 1; i >= 0; i-- {
		migration := g.migrations[i]
		if g.migrationDidRun(migration) {
			return migration, nil
		}
	}
	return nil, ErrNoRunMigration
}

// RollbackMigration undo a migration.
func (g *Gormigrate) RollbackMigration(m *Migration) error {
	g.begin()
	defer g.rollback()

	if err := g.rollbackMigration(m); err != nil {
		return err
	}
	return g.commit()
}

func (g *Gormigrate) rollbackMigration(m *Migration) error {
	if m.Rollback == nil {
		return ErrRollbackImpossible
	}

	if err := m.Rollback(g.tx); err != nil {
		return err
	}

	sql := fmt.Sprintf("DELETE FROM %s WHERE %s = ?", g.options.TableName, g.options.IDColumnName)
	if err := g.tx.Exec(sql, m.ID).Error; err != nil {
		return err
	}
	return nil
}

func (g *Gormigrate) runInitSchema() error {
	if err := g.initSchema(g.tx); err != nil {
		return err
	}
	if err := g.insertMigration(initSchemaMigrationId); err != nil {
		return err
	}

	for _, migration := range g.migrations {
		if err := g.insertMigration(migration.ID); err != nil {
			return err
		}
	}

	return nil
}

func (g *Gormigrate) runMigration(migration *Migration) error {
	if len(migration.ID) == 0 {
		return ErrMissingID
	}

	if !g.migrationDidRun(migration) {
		if err := migration.Migrate(g.tx); err != nil {
			return err
		}

		if err := g.insertMigration(migration.ID); err != nil {
			return err
		}
	}
	return nil
}

func (g *Gormigrate) createMigrationTableIfNotExists() error {
	if g.tx.HasTable(g.options.TableName) {
		return nil
	}

	sql := fmt.Sprintf("CREATE TABLE %s (%s VARCHAR(%d) PRIMARY KEY)", g.options.TableName, g.options.IDColumnName, g.options.IDColumnSize)
	if err := g.tx.Exec(sql).Error; err != nil {
		return err
	}
	return nil
}

func (g *Gormigrate) migrationDidRun(m *Migration) bool {
	var count int
	g.tx.
		Table(g.options.TableName).
		Where(fmt.Sprintf("%s = ?", g.options.IDColumnName), m.ID).
		Count(&count)
	return count > 0
}

// The schema can be initialised only if it hasn't been initialised yet
// and no other migration has been applied already.
func (g *Gormigrate) canInitializeSchema() bool {
	if g.migrationDidRun(&Migration{ID: initSchemaMigrationId}) {
		return false
	}

	// If the ID doesn't exist, we also want the list of migrations to be empty
	var count int
	g.tx.
		Table(g.options.TableName).
		Count(&count)
	return count == 0
}

func (g *Gormigrate) insertMigration(id string) error {
	sql := fmt.Sprintf("INSERT INTO %s (%s) VALUES (?)", g.options.TableName, g.options.IDColumnName)
	return g.tx.Exec(sql, id).Error
}

func (g *Gormigrate) begin() {
	if g.options.UseTransaction {
		g.tx = g.db.Begin()
	} else {
		g.tx = g.db
	}
}

func (g *Gormigrate) commit() error {
	if g.options.UseTransaction {
		if err := g.tx.Commit().Error; err != nil {
			return err
		}
	}
	return nil
}

func (g *Gormigrate) rollback() {
	if g.options.UseTransaction {
		g.tx.Rollback()
	}
}
<commit_msg>Small code style improvements<commit_after>package gormigrate

import (
	"errors"
	"fmt"

	"github.com/jinzhu/gorm"
)

const (
	initSchemaMigrationID = "SCHEMA_INIT"
)

// MigrateFunc is the func signature for migrating.
type MigrateFunc func(*gorm.DB) error

// RollbackFunc is the func signature for rollbacking.
type RollbackFunc func(*gorm.DB) error

// InitSchemaFunc is the func signature for initializing the schema.
type InitSchemaFunc func(*gorm.DB) error

// Options define options for all migrations.
type Options struct {
	// TableName is the migration table.
	TableName string
	// IDColumnName is the name of column where the migration id will be stored.
	IDColumnName string
	// IDColumnSize is the length of the migration id column
	IDColumnSize int
	// UseTransaction makes Gormigrate execute migrations inside a single transaction.
	// Keep in mind that not all databases support DDL commands inside transactions.
	UseTransaction bool
}

// Migration represents a database migration (a modification to be made on the database).
type Migration struct {
	// ID is the migration identifier. Usually a timestamp like "201601021504".
	ID string
	// Migrate is a function that will br executed while running this migration.
	Migrate MigrateFunc
	// Rollback will be executed on rollback. Can be nil.
	Rollback RollbackFunc
}

// Gormigrate represents a collection of all migrations of a database schema.
type Gormigrate struct {
	db         *gorm.DB
	tx         *gorm.DB
	options    *Options
	migrations []*Migration
	initSchema InitSchemaFunc
}

// ReservedIDError is returned when a migration is using a reserved ID
type ReservedIDError struct {
	ID string
}

func (e *ReservedIDError) Error() string {
	return fmt.Sprintf(`gormigrate: Reserved migration ID: "%s"`, e.ID)
}

// DuplicatedIDError is returned when more than one migration have the same ID
type DuplicatedIDError struct {
	ID string
}

func (e *DuplicatedIDError) Error() string {
	return fmt.Sprintf(`gormigrate: Duplicated migration ID: "%s"`, e.ID)
}

var (
	// DefaultOptions can be used if you don't want to think about options.
	DefaultOptions = &Options{
		TableName:      "migrations",
		IDColumnName:   "id",
		IDColumnSize:   255,
		UseTransaction: false,
	}

	// ErrRollbackImpossible is returned when trying to rollback a migration
	// that has no rollback function.
	ErrRollbackImpossible = errors.New("gormigrate: It's impossible to rollback this migration")

	// ErrNoMigrationDefined is returned when no migration is defined.
	ErrNoMigrationDefined = errors.New("gormigrate: No migration defined")

	// ErrMissingID is returned when the ID od migration is equal to ""
	ErrMissingID = errors.New("gormigrate: Missing ID in migration")

	// ErrNoRunMigration is returned when any run migration was found while
	// running RollbackLast
	ErrNoRunMigration = errors.New("gormigrate: Could not find last run migration")

	// ErrMigrationIDDoesNotExist is returned when migrating or rolling back to a migration ID that
	// does not exist in the list of migrations
	ErrMigrationIDDoesNotExist = errors.New("gormigrate: Tried to migrate to an ID that doesn't exist")
)

// New returns a new Gormigrate.
func New(db *gorm.DB, options *Options, migrations []*Migration) *Gormigrate {
	if options.TableName == "" {
		options.TableName = DefaultOptions.TableName
	}
	if options.IDColumnName == "" {
		options.IDColumnName = DefaultOptions.IDColumnName
	}
	if options.IDColumnSize == 0 {
		options.IDColumnSize = DefaultOptions.IDColumnSize
	}
	return &Gormigrate{
		db:         db,
		options:    options,
		migrations: migrations,
	}
}

// InitSchema sets a function that is run if no migration is found.
// The idea is preventing to run all migrations when a new clean database
// is being migrating. In this function you should create all tables and
// foreign key necessary to your application.
func (g *Gormigrate) InitSchema(initSchema InitSchemaFunc) {
	g.initSchema = initSchema
}

// Migrate executes all migrations that did not run yet.
func (g *Gormigrate) Migrate() error {
	if !g.hasMigrations() {
		return ErrNoMigrationDefined
	}
	var targetMigrationID string
	if len(g.migrations) > 0 {
		targetMigrationID = g.migrations[len(g.migrations)-1].ID
	}
	return g.migrate(targetMigrationID)
}

// MigrateTo executes all migrations that did not run yet up to the migration that matches `migrationID`.
func (g *Gormigrate) MigrateTo(migrationID string) error {
	if err := g.checkIDExist(migrationID); err != nil {
		return err
	}
	return g.migrate(migrationID)
}

func (g *Gormigrate) migrate(migrationID string) error {
	if !g.hasMigrations() {
		return ErrNoMigrationDefined
	}

	if err := g.checkReservedID(); err != nil {
		return err
	}

	if err := g.checkDuplicatedID(); err != nil {
		return err
	}

	g.begin()
	defer g.rollback()

	if err := g.createMigrationTableIfNotExists(); err != nil {
		return err
	}

	if g.initSchema != nil && g.canInitializeSchema() {
		if err := g.runInitSchema(); err != nil {
			return err
		}
		return g.commit()
	}

	for _, migration := range g.migrations {
		if err := g.runMigration(migration); err != nil {
			return err
		}
		if migrationID != "" && migration.ID == migrationID {
			break
		}
	}
	return g.commit()
}

// There are migrations to apply if either there's a defined
// initSchema function or if the list of migrations is not empty.
func (g *Gormigrate) hasMigrations() bool {
	return g.initSchema != nil || len(g.migrations) > 0
}

// Check whether any migration is using a reserved ID.
// For now there's only have one reserved ID, but there may be more in the future.
func (g *Gormigrate) checkReservedID() error {
	for _, m := range g.migrations {
		if m.ID == initSchemaMigrationID {
			return &ReservedIDError{ID: m.ID}
		}
	}
	return nil
}

func (g *Gormigrate) checkDuplicatedID() error {
	lookup := make(map[string]struct{}, len(g.migrations))
	for _, m := range g.migrations {
		if _, ok := lookup[m.ID]; ok {
			return &DuplicatedIDError{ID: m.ID}
		}
		lookup[m.ID] = struct{}{}
	}
	return nil
}

func (g *Gormigrate) checkIDExist(migrationID string) error {
	for _, migrate := range g.migrations {
		if migrate.ID == migrationID {
			return nil
		}
	}
	return ErrMigrationIDDoesNotExist
}

// RollbackLast undo the last migration
func (g *Gormigrate) RollbackLast() error {
	if len(g.migrations) == 0 {
		return ErrNoMigrationDefined
	}

	g.begin()
	defer g.rollback()

	lastRunMigration, err := g.getLastRunMigration()
	if err != nil {
		return err
	}

	if err := g.rollbackMigration(lastRunMigration); err != nil {
		return err
	}
	return g.commit()
}

// RollbackTo undoes migrations up to the given migration that matches the `migrationID`.
// Migration with the matching `migrationID` is not rolled back.
func (g *Gormigrate) RollbackTo(migrationID string) error {
	if len(g.migrations) == 0 {
		return ErrNoMigrationDefined
	}

	if err := g.checkIDExist(migrationID); err != nil {
		return err
	}

	g.begin()
	defer g.rollback()

	for i := len(g.migrations) - 1; i >= 0; i-- {
		migration := g.migrations[i]
		if migration.ID == migrationID {
			break
		}
		if g.migrationDidRun(migration) {
			if err := g.rollbackMigration(migration); err != nil {
				return err
			}
		}
	}
	return g.commit()
}

func (g *Gormigrate) getLastRunMigration() (*Migration, error) {
	for i := len(g.migrations) - 1; i >= 0; i-- {
		migration := g.migrations[i]
		if g.migrationDidRun(migration) {
			return migration, nil
		}
	}
	return nil, ErrNoRunMigration
}

// RollbackMigration undo a migration.
func (g *Gormigrate) RollbackMigration(m *Migration) error {
	g.begin()
	defer g.rollback()

	if err := g.rollbackMigration(m); err != nil {
		return err
	}
	return g.commit()
}

func (g *Gormigrate) rollbackMigration(m *Migration) error {
	if m.Rollback == nil {
		return ErrRollbackImpossible
	}

	if err := m.Rollback(g.tx); err != nil {
		return err
	}

	sql := fmt.Sprintf("DELETE FROM %s WHERE %s = ?", g.options.TableName, g.options.IDColumnName)
	return g.tx.Exec(sql, m.ID).Error
}

func (g *Gormigrate) runInitSchema() error {
	if err := g.initSchema(g.tx); err != nil {
		return err
	}
	if err := g.insertMigration(initSchemaMigrationID); err != nil {
		return err
	}

	for _, migration := range g.migrations {
		if err := g.insertMigration(migration.ID); err != nil {
			return err
		}
	}

	return nil
}

func (g *Gormigrate) runMigration(migration *Migration) error {
	if len(migration.ID) == 0 {
		return ErrMissingID
	}

	if !g.migrationDidRun(migration) {
		if err := migration.Migrate(g.tx); err != nil {
			return err
		}

		if err := g.insertMigration(migration.ID); err != nil {
			return err
		}
	}
	return nil
}

func (g *Gormigrate) createMigrationTableIfNotExists() error {
	if g.tx.HasTable(g.options.TableName) {
		return nil
	}

	sql := fmt.Sprintf("CREATE TABLE %s (%s VARCHAR(%d) PRIMARY KEY)", g.options.TableName, g.options.IDColumnName, g.options.IDColumnSize)
	return g.tx.Exec(sql).Error
}

func (g *Gormigrate) migrationDidRun(m *Migration) bool {
	var count int
	g.tx.
		Table(g.options.TableName).
		Where(fmt.Sprintf("%s = ?", g.options.IDColumnName), m.ID).
		Count(&count)
	return count > 0
}

// The schema can be initialised only if it hasn't been initialised yet
// and no other migration has been applied already.
func (g *Gormigrate) canInitializeSchema() bool {
	if g.migrationDidRun(&Migration{ID: initSchemaMigrationID}) {
		return false
	}

	// If the ID doesn't exist, we also want the list of migrations to be empty
	var count int
	g.tx.
		Table(g.options.TableName).
		Count(&count)
	return count == 0
}

func (g *Gormigrate) insertMigration(id string) error {
	sql := fmt.Sprintf("INSERT INTO %s (%s) VALUES (?)", g.options.TableName, g.options.IDColumnName)
	return g.tx.Exec(sql, id).Error
}

func (g *Gormigrate) begin() {
	if g.options.UseTransaction {
		g.tx = g.db.Begin()
	} else {
		g.tx = g.db
	}
}

func (g *Gormigrate) commit() error {
	if g.options.UseTransaction {
		return g.tx.Commit().Error
	}
	return nil
}

func (g *Gormigrate) rollback() {
	if g.options.UseTransaction {
		g.tx.Rollback()
	}
}
<|endoftext|>


<commit_before>package gossh_test

import (
	. "github.com/dcapwell/gossh"
	"github.com/dcapwell/gossh/workpool"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"log"
)

var _ = Describe("Gossh", func() {
  Describe("multi hosts", func() {
    ssh := NewSsh()

    Describe("run multiple ssh commands concurrently", func() {
      Context("with one host", func() {
        Context("and valid command", func() {
          rsp, err := ssh.Run([]string{"localhost"}, "date", Options{})
          data := chanToSlize(rsp.Responses)

          It("should succeed", func() {
            Expect(err).To(BeNil())
          })

          It("with one result", func() {
            Expect(data).To(HaveLen(1))
            Expect(data[0].Hostname).To(Equal("localhost"))
            Expect(data[0].Response.Code).To(Equal(workpool.SUCCESS))
            Expect(data[0].Response.Stdout).To(ContainSubstring("PST"))
          })
        })
        Context("with invalid command", func() {
          rsp, err := ssh.Run([]string{"localhost"}, "thiscmdreallyshouldntexist", Options{})
          data := chanToSlize(rsp.Responses)

          It("should not return error", func() {
            Expect(err).To(BeNil())
          })
          It("should have only one response", func() {
            Expect(data).To(HaveLen(1))
          })
          It("should have response from host", func() {
            localRsp := data[0]
            Expect(localRsp.Hostname).To(Equal("localhost"))
          })
          It("should have failed", func() {
            localRsp := data[0]
            Expect(localRsp.Response).ShouldNot(BeNil())
            Expect(localRsp.Response.Code).To(BeGreaterThan(0))
          })
        })
      })
      Context("with mulitple hosts", func() {
        Context("with valid command", func() {
          rsp, err := ssh.Run([]string{"localhost", "localhost"}, "date", Options{})
          data := chanToSlize(rsp.Responses)

          It("should not return error", func() {
            Expect(err).To(BeNil())
          })
          It("should have only two response", func() {
            Expect(data).To(HaveLen(2))
          })
          It("should have response from host", func() {
            localRsp := data[0]
            Expect(localRsp.Hostname).To(Equal("localhost"))
            localRsp = data[1]
            Expect(localRsp.Hostname).To(Equal("localhost"))
          })
          It("should have successed", func() {
            Expect(data[0].Response.Code).To(Equal(workpool.SUCCESS))
            Expect(data[0].Response.Stdout).To(ContainSubstring("PST"))

            Expect(data[1].Response.Code).To(Equal(workpool.SUCCESS))
            Expect(data[1].Response.Stdout).To(ContainSubstring("PST"))
          })
        })
        Context("with invalid command", func() {
          rsp, err := ssh.Run([]string{"localhost", "localhost"}, "thiscmdreallyshouldntexist", Options{})
          data := chanToSlize(rsp.Responses)

          log.Printf("Response from multi requests, single host, invalid cmd: %v\n", rsp)

          It("should not return error", func() {
            Expect(err).To(BeNil())
          })
          It("should have only two response", func() {
            Expect(data).To(HaveLen(2))
          })
          It("should have response from host", func() {
            Expect(data[0].Hostname).To(Equal("localhost"))
            Expect(data[1].Hostname).To(Equal("localhost"))
          })
          It("should have failed", func() {
            localRsp := data[0]
            Expect(localRsp.Response).ShouldNot(BeNil())
            Expect(localRsp.Response.Code).To(BeGreaterThan(0))

            localRsp = data[1]
            Expect(localRsp.Response).ShouldNot(BeNil())
            Expect(localRsp.Response.Code).To(BeGreaterThan(0))
          })
        })
      })
    })
  })
})

func chanToSlize(ch chan SshResponseContext) []SshResponseContext {
	data := make([]SshResponseContext, 0)
	for cxt := range ch {
		data = append(data, cxt)
	}
	return data
}
<commit_msg>tests now use UTC so output doesnt depend on location<commit_after>package gossh_test

import (
	. "github.com/dcapwell/gossh"
	"github.com/dcapwell/gossh/workpool"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"log"
)

var _ = Describe("Gossh", func() {
  Describe("multi hosts", func() {
    ssh := NewSsh()

    Describe("run multiple ssh commands concurrently", func() {
      Context("with one host", func() {
        Context("and valid command", func() {
          rsp, err := ssh.Run([]string{"localhost"}, "date -u", Options{})
          data := chanToSlize(rsp.Responses)

          It("should succeed", func() {
            Expect(err).To(BeNil())
          })

          It("with one result", func() {
            Expect(data).To(HaveLen(1))
            Expect(data[0].Hostname).To(Equal("localhost"))
            Expect(data[0].Response.Code).To(Equal(workpool.SUCCESS))
            Expect(data[0].Response.Stdout).To(ContainSubstring("UTC"))
          })
        })
        Context("with invalid command", func() {
          rsp, err := ssh.Run([]string{"localhost"}, "thiscmdreallyshouldntexist", Options{})
          data := chanToSlize(rsp.Responses)

          It("should not return error", func() {
            Expect(err).To(BeNil())
          })
          It("should have only one response", func() {
            Expect(data).To(HaveLen(1))
          })
          It("should have response from host", func() {
            localRsp := data[0]
            Expect(localRsp.Hostname).To(Equal("localhost"))
          })
          It("should have failed", func() {
            localRsp := data[0]
            Expect(localRsp.Response).ShouldNot(BeNil())
            Expect(localRsp.Response.Code).To(BeGreaterThan(0))
          })
        })
      })
      Context("with mulitple hosts", func() {
        Context("with valid command", func() {
          rsp, err := ssh.Run([]string{"localhost", "localhost"}, "date -u", Options{})
          data := chanToSlize(rsp.Responses)

          It("should not return error", func() {
            Expect(err).To(BeNil())
          })
          It("should have only two response", func() {
            Expect(data).To(HaveLen(2))
          })
          It("should have response from host", func() {
            localRsp := data[0]
            Expect(localRsp.Hostname).To(Equal("localhost"))
            localRsp = data[1]
            Expect(localRsp.Hostname).To(Equal("localhost"))
          })
          It("should have successed", func() {
            Expect(data[0].Response.Code).To(Equal(workpool.SUCCESS))
            Expect(data[0].Response.Stdout).To(ContainSubstring("UTC"))

            Expect(data[1].Response.Code).To(Equal(workpool.SUCCESS))
            Expect(data[1].Response.Stdout).To(ContainSubstring("UTC"))
          })
        })
        Context("with invalid command", func() {
          rsp, err := ssh.Run([]string{"localhost", "localhost"}, "thiscmdreallyshouldntexist", Options{})
          data := chanToSlize(rsp.Responses)

          log.Printf("Response from multi requests, single host, invalid cmd: %v\n", rsp)

          It("should not return error", func() {
            Expect(err).To(BeNil())
          })
          It("should have only two response", func() {
            Expect(data).To(HaveLen(2))
          })
          It("should have response from host", func() {
            Expect(data[0].Hostname).To(Equal("localhost"))
            Expect(data[1].Hostname).To(Equal("localhost"))
          })
          It("should have failed", func() {
            localRsp := data[0]
            Expect(localRsp.Response).ShouldNot(BeNil())
            Expect(localRsp.Response.Code).To(BeGreaterThan(0))

            localRsp = data[1]
            Expect(localRsp.Response).ShouldNot(BeNil())
            Expect(localRsp.Response.Code).To(BeGreaterThan(0))
          })
        })
      })
    })
  })
})

func chanToSlize(ch chan SshResponseContext) []SshResponseContext {
	data := make([]SshResponseContext, 0)
	for cxt := range ch {
		data = append(data, cxt)
	}
	return data
}
<|endoftext|>


<commit_before>/*
 * Copyright (c) 2017 AlexRuzin (stan.ruzin@gmail.com)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package gofs

import (
    "testing"
    "time"
)

func TestIOSanity(t *testing.T) {
	out("[+] Running Standard I/O Sanity Test...")

    var header = create_db("<path>")
    if header == nil {
        drive_fail("TEST1: Failed to obtain header", t)
    }
	out("[+] Test 1 PASS")

    // The root file "/" must at least exist
    if file, err := header.create("/"); file != nil && err != STATUS_ERROR {
        drive_fail("TEST2: Failed to return root handle", t)
    }
	out("[+] Test 2 PASS")

    /*
     * Try to delete the root file "/"
     */
    if header.delete("/") == 0 {
        drive_fail("TEST3: Cannot delete root -- critical", t)
    }
	out("[+] Test 3 PASS")

    /*
     * Attempt to write to a nonexistant file
     */
	var data = []byte{ 1, 2 }
    if header.write("/folder5/folder5/file5", data) != STATUS_ERROR {
        drive_fail("TEST4: Cannot write to a nonexistant file", t)
    }
	out("[+] Test 4 PASS")

    /*
     * Attempt to create a new file0
     */
    if file, err := header.create("/folder0/folder0/file0"); file == nil || err == STATUS_ERROR {
        drive_fail("TEST5.0: file0 cannot be created", t)
    }
	out("[+] Test 5.0 PASS")

    /*
     * Attempt to create a new file0, this will fail since it should already exist
     */
    if file, err := header.create("/folder0/folder0/file0"); file != nil && err != STATUS_EXISTS {
        drive_fail("TEST5.1: file0 cannot be created twice", t)
    }
	out("[+] Test 5.1 PASS")

    /*
     * Write some data into file0
     */
    data = []byte{ 1, 2, 3, 4 }
    if header.write("/folder0/folder0/file0", data) != STATUS_OK {
        drive_fail("TEST6: Failed to write data in file0", t)
    }
	out("[+] Test 6 PASS")

    /*
     * Attempt to create a new file3
     */
    if file, err := header.create("/folder1/folder0/file3"); file == nil || err == STATUS_ERROR {
        drive_fail("TEST7: file3 cannot be created", t)
    }
	out("[+] Test 7 PASS")

    /*
     * Write some data into file3
     */
    var data2 = []byte{ 1, 2, 3, 4, 5, 6, 7 }
    if header.write("/folder1/folder0/file3", data2) != 0 {
        drive_fail("TEST8: Failed to write data in file3", t)
    }
	out("[+] Test 8 PASS")

    /*
     * Write some data into file3
     */
    if header.write("/folder1/folder0/file3", data2) != 0 {
        drive_fail("TEST8.1: Failed to write data in file3", t)
    }
	out("[+] Test 8.1 PASS")

    /*
     * Read the written data from file0 and compare
     */
    output_data := header.read("/folder0/folder0/file0")
    if output_data == nil || len(output_data) != len(data) || header.t_size - 7 /* len(file3) */ != uint(len(data)) {
        drive_fail("TEST9: Failed to read data from file0", t)
    }
	out("[+] Test 9 PASS")

    /*
     * Read the written data from file3 and compare
     */
    output_data = header.read("/folder1/folder0/file3")
    if output_data == nil || len(output_data) != len(data2) || header.t_size - 4 /* len(file0) */ != uint(len(data2)) {
        drive_fail("TEST10: Failed to read data from file3", t)
    }
	out("[+] Test 10 PASS")

    /*
     * Write other data to file0
     */
    data = []byte{ 1, 2, 3 }
    if header.write("/folder0/folder0/file0", data) != 0 {
        drive_fail("TEST11: Failed to write data in file1", t)
    }
	out("[+] Test 11 PASS")

    /*
     * Read the new data from file0
     */
    output_data = header.read("/folder0/folder0/file0")
    if output_data == nil || len(output_data) != len(data) {
        drive_fail("TEST12: Failed to read data from file1", t)
    }
	out("[+] Test 12 PASS")

    /*
     * Attempt to create a new file5. This will be a blank file
     */
    if file, err := header.create("/folder2/file5"); file == nil || err == STATUS_ERROR {
        drive_fail("TEST13: file3 cannot be created", t)
    }
	out("[+] Test 13 PASS")

    /*
     * Delete file0 -- complete this
     */

	/*
	 * Create just a folder
	 */
    if file, err := header.create("/folder2/file5/"); file == nil || err == STATUS_ERROR {
        drive_fail("TEST15: folder file5 cannot be created", t)
    }
	out("[+] Test 15 PASS")

     /*
      * Unmount/commit database to file
      */
    if header.unmount_db(nil) != 0 {
        drive_fail("TEST16: Failed to commit database", t)
    }
	out("[+] Test 16 PASS")

    time.Sleep(10000)
}

func drive_fail(output string, t *testing.T) {
	t.Errorf(output)
	t.FailNow()
}
<commit_msg>Fixed the major sync problem that was breaking the I/O engine.<commit_after>/*
 * Copyright (c) 2017 AlexRuzin (stan.ruzin@gmail.com)
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

package gofs

import (
    "testing"
    "time"
)

func TestIOSanity(t *testing.T) {
	out("[+] Running Standard I/O Sanity Test...")

    var header = create_db("<path>")
    if header == nil {
        drive_fail("TEST1: Failed to obtain header", t)
    }
	out("[+] Test 1 PASS")

    // The root file "/" must at least exist
    if file, err := header.create("/"); file != nil && err != STATUS_ERROR {
        drive_fail("TEST2: Failed to return root handle", t)
    }
	out("[+] Test 2 PASS")

    /*
     * Try to delete the root file "/"
     */
    if header.delete("/") == 0 {
        drive_fail("TEST3: Cannot delete root -- critical", t)
    }
	out("[+] Test 3 PASS")

    /*
     * Attempt to write to a nonexistant file
     */
	var data = []byte{ 1, 2 }
    if header.write("/folder5/folder5/file5", data) != STATUS_ERROR {
        drive_fail("TEST4: Cannot write to a nonexistant file", t)
    }
	out("[+] Test 4 PASS")

    /*
     * Attempt to create a new file0
     */
    if file, err := header.create("/folder0/folder0/file0"); file == nil || err == STATUS_ERROR {
        drive_fail("TEST5.0: file0 cannot be created", t)
    }
	out("[+] Test 5.0 PASS")

    /*
     * Attempt to create a new file0, this will fail since it should already exist
     */
    if file, err := header.create("/folder0/folder0/file0"); file != nil && err != STATUS_EXISTS {
        drive_fail("TEST5.1: file0 cannot be created twice", t)
    }
	out("[+] Test 5.1 PASS")

    /*
     * Write some data into file0
     */
    data = []byte{ 1, 2, 3, 4 }
    if header.write("/folder0/folder0/file0", data) != STATUS_OK {
        drive_fail("TEST6: Failed to write data in file0", t)
    }
	out("[+] Test 6 PASS")

    /*
     * Attempt to create a new file3
     */
    if file, err := header.create("/folder1/folder0/file3"); file == nil || err == STATUS_ERROR {
        drive_fail("TEST7: file3 cannot be created", t)
    }
	out("[+] Test 7 PASS")

    /*
     * Write some data into file3
     */
    var data2 = []byte{ 1, 2, 3, 4, 5, 6, 7 }
    if header.write("/folder1/folder0/file3", data2) != 0 {
        drive_fail("TEST8: Failed to write data in file3", t)
    }
	out("[+] Test 8 PASS")

    /*
     * Write some data into file3
     */
    if header.write("/folder1/folder0/file3", data2) != 0 {
        drive_fail("TEST8.1: Failed to write data in file3", t)
    }
	out("[+] Test 8.1 PASS")

    /*
     * Read the written data from file0 and compare
     */
    output_data := header.read("/folder0/folder0/file0")
    if output_data == nil || len(output_data) != len(data) || header.t_size - 7 /* len(file3) */ != uint(len(data)) {
        drive_fail("TEST9: Failed to read data from file0", t)
    }
	out("[+] Test 9 PASS")

    /*
     * Read the written data from file3 and compare
     */
    output_data = header.read("/folder1/folder0/file3")
    if output_data == nil || len(output_data) != len(data2) || header.t_size - 4 /* len(file0) */ != uint(len(data2)) {
        drive_fail("TEST10: Failed to read data from file3", t)
    }
	out("[+] Test 10 PASS")

    /*
     * Write other data to file0
     */
    data = []byte{ 1, 2, 3 }
    if header.write("/folder0/folder0/file0", data) != 0 {
        drive_fail("TEST11: Failed to write data in file1", t)
    }
	out("[+] Test 11 PASS")

    /*
     * Read the new data from file0
     */
    output_data = header.read("/folder0/folder0/file0")
    if output_data == nil || len(output_data) != len(data) {
        drive_fail("TEST12: Failed to read data from file1", t)
    }
	out("[+] Test 12 PASS")

    /*
     * Attempt to create a new file5. This will be a blank file
     */
    if file, err := header.create("/folder2/file7"); file == nil || err == STATUS_ERROR {
        drive_fail("TEST13: file3 cannot be created", t)
    }
	out("[+] Test 13 PASS")

    /*
     * Delete file0 -- complete this
     */
	// FIXME/ADDME

	/*
	 * Create just a folder
	 */
    if file, err := header.create("/folder2/file5/"); file == nil || err == STATUS_ERROR {
        drive_fail("TEST15: folder file5 cannot be created", t)
    }
	out("[+] Test 15 PASS")

     /*
      * Unmount/commit database to file
      */
    if header.unmount_db(nil) != 0 {
        drive_fail("TEST16: Failed to commit database", t)
    }
	out("[+] Test 16 PASS")

    time.Sleep(10000)
}

func drive_fail(output string, t *testing.T) {
	t.Errorf(output)
	t.FailNow()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"time"

	. "github.com/russross/codegrinder/types"
	"github.com/spf13/cobra"
)

func CommandSave(cmd *cobra.Command, args []string) {
	mustLoadConfig()
	now := time.Now()

	// find the directory
	dir := ""
	switch len(args) {
	case 0:
		dir = "."
	case 1:
		dir = args[0]
	default:
		cmd.Help()
		return
	}

	problem, _, commit, _ := gather(now, dir)
	commit.Action = ""
	commit.Note = "saving from grind tool"

	// send the commit to the server
	signed := new(Commit)
	mustPostObject(fmt.Sprintf("/assignments/%d/commits", commit.AssignmentID), nil, commit, signed)
	log.Printf("problem %s step %d saved", problem.Unique, commit.Step)
}

func gather(now time.Time, startDir string) (*Problem, *Assignment, *Commit, *DotFileInfo) {
	// find the .grind file containing the problem set info
	dotfile, problemSetDir, problemDir := findDotFile(startDir)

	// get the assignment
	assignment := new(Assignment)
	mustGetObject(fmt.Sprintf("/assignments/%d", dotfile.AssignmentID), nil, assignment)

	// get the problem
	unique := ""
	if len(dotfile.Problems) == 1 {
		// only one problem? files should be in dotfile directory
		for u := range dotfile.Problems {
			unique = u
		}
		problemDir = problemSetDir
	} else {
		// use the subdirectory name to identify the problem
		if problemDir == "" {
			log.Printf("you must identify the problem within this problem set")
			log.Printf("  either run this from with the problem directory, or")
			log.Fatalf("  identify it as a parameter in the command")
		}
		_, unique = filepath.Split(problemDir)
	}
	info := dotfile.Problems[unique]
	if info == nil {
		log.Fatalf("unable to recognize the problem based on the directory name of %q", unique)
	}
	problem := new(Problem)
	mustGetObject(fmt.Sprintf("/problems/%d", info.ID), nil, problem)

	// gather the commit files from the file system
	files := make(map[string]string)
	err := filepath.Walk(problemDir, func(path string, stat os.FileInfo, err error) error {
		// skip errors, directories, non-regular files
		if err != nil {
			return err
		}
		if path == problemDir {
			// descent into the main directory
			return nil
		}
		if stat.IsDir() {
			return filepath.SkipDir
		}
		if !stat.Mode().IsRegular() {
			return nil
		}
		_, name := filepath.Split(path)

		// skip our config file
		if name == perProblemSetDotFile {
			return nil
		}

		if info.Whitelist[name] {
			contents, err := ioutil.ReadFile(path)
			if err != nil {
				return err
			}
			files[name] = string(contents)
		} else {
			log.Printf("skipping %q which is not a file introduced by the problem", name)
		}
		return nil
	})
	if err != nil {
		log.Fatalf("walk error: %v", err)
	}
	if len(files) != len(info.Whitelist) {
		log.Printf("did not find all the expected files")
		for name := range info.Whitelist {
			if _, ok := files[name]; !ok {
				log.Printf("  %s not found", name)
			}
		}
		log.Fatalf("all expected files must be present")
	}

	// form a commit object
	commit := &Commit{
		ID:           0,
		AssignmentID: dotfile.AssignmentID,
		ProblemID:    info.ID,
		Step:         info.Step,
		Files:        files,
		CreatedAt:    now,
		UpdatedAt:    now,
	}

	return problem, assignment, commit, dotfile
}

func findDotFile(startDir string) (dotfile *DotFileInfo, problemSetDir, problemDir string) {
	abs := false
	problemSetDir, problemDir = startDir, ""
	for {
		path := filepath.Join(problemSetDir, perProblemSetDotFile)
		if _, err := os.Stat(path); err != nil {
			if os.IsNotExist(err) {
				if !abs {
					abs = true
					path, err := filepath.Abs(problemSetDir)
					if err != nil {
						log.Fatalf("error finding absolute path of %s: %v", problemSetDir, err)
					}
					problemSetDir = path
				}

				// try moving up a directory
				problemDir = problemSetDir
				problemSetDir = filepath.Dir(problemSetDir)
				if problemSetDir == problemDir {
					log.Fatalf("unable to find %s in %s or an ancestor directory", perProblemSetDotFile, startDir)
				}
				log.Printf("could not find %s in %s, trying %s", perProblemSetDotFile, problemDir, problemSetDir)
				continue
			}

			log.Fatalf("error searching for %s in %s: %v", perProblemSetDotFile, problemSetDir, err)
		}
		break
	}

	// read the .grind file
	path := filepath.Join(problemSetDir, perProblemSetDotFile)
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		log.Fatalf("error reading %s: %v", path, err)
	}
	dotfile = new(DotFileInfo)
	if err := json.Unmarshal(contents, dotfile); err != nil {
		log.Fatalf("error parsing %s: %v", path, err)
	}
	dotfile.Path = path

	return dotfile, problemSetDir, problemDir
}
<commit_msg>saving works from grind<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"time"

	. "github.com/russross/codegrinder/types"
	"github.com/spf13/cobra"
)

func CommandSave(cmd *cobra.Command, args []string) {
	mustLoadConfig()
	now := time.Now()

	// find the directory
	dir := ""
	switch len(args) {
	case 0:
		dir = "."
	case 1:
		dir = args[0]
	default:
		cmd.Help()
		return
	}

	problem, _, commit, _ := gather(now, dir)
	commit.Action = ""
	commit.Note = "saving from grind tool"
	unsigned := &CommitBundle{Commit: commit}

	// send the commit to the server
	signed := new(CommitBundle)
	mustPostObject("/commit_bundles/unsigned", nil, unsigned, signed)
	log.Printf("problem %s step %d saved", problem.Unique, commit.Step)
}

func gather(now time.Time, startDir string) (*Problem, *Assignment, *Commit, *DotFileInfo) {
	// find the .grind file containing the problem set info
	dotfile, problemSetDir, problemDir := findDotFile(startDir)

	// get the assignment
	assignment := new(Assignment)
	mustGetObject(fmt.Sprintf("/assignments/%d", dotfile.AssignmentID), nil, assignment)

	// get the problem
	unique := ""
	if len(dotfile.Problems) == 1 {
		// only one problem? files should be in dotfile directory
		for u := range dotfile.Problems {
			unique = u
		}
		problemDir = problemSetDir
	} else {
		// use the subdirectory name to identify the problem
		if problemDir == "" {
			log.Printf("you must identify the problem within this problem set")
			log.Printf("  either run this from with the problem directory, or")
			log.Fatalf("  identify it as a parameter in the command")
		}
		_, unique = filepath.Split(problemDir)
	}
	info := dotfile.Problems[unique]
	if info == nil {
		log.Fatalf("unable to recognize the problem based on the directory name of %q", unique)
	}
	problem := new(Problem)
	mustGetObject(fmt.Sprintf("/problems/%d", info.ID), nil, problem)

	// gather the commit files from the file system
	files := make(map[string]string)
	err := filepath.Walk(problemDir, func(path string, stat os.FileInfo, err error) error {
		// skip errors, directories, non-regular files
		if err != nil {
			return err
		}
		if path == problemDir {
			// descent into the main directory
			return nil
		}
		if stat.IsDir() {
			return filepath.SkipDir
		}
		if !stat.Mode().IsRegular() {
			return nil
		}
		_, name := filepath.Split(path)

		// skip our config file
		if name == perProblemSetDotFile {
			return nil
		}

		if info.Whitelist[name] {
			contents, err := ioutil.ReadFile(path)
			if err != nil {
				return err
			}
			files[name] = string(contents)
		} else {
			log.Printf("skipping %q which is not a file introduced by the problem", name)
		}
		return nil
	})
	if err != nil {
		log.Fatalf("walk error: %v", err)
	}
	if len(files) != len(info.Whitelist) {
		log.Printf("did not find all the expected files")
		for name := range info.Whitelist {
			if _, ok := files[name]; !ok {
				log.Printf("  %s not found", name)
			}
		}
		log.Fatalf("all expected files must be present")
	}

	// form a commit object
	commit := &Commit{
		ID:           0,
		AssignmentID: dotfile.AssignmentID,
		ProblemID:    info.ID,
		Step:         info.Step,
		Files:        files,
		CreatedAt:    now,
		UpdatedAt:    now,
	}

	return problem, assignment, commit, dotfile
}

func findDotFile(startDir string) (dotfile *DotFileInfo, problemSetDir, problemDir string) {
	abs := false
	problemSetDir, problemDir = startDir, ""
	for {
		path := filepath.Join(problemSetDir, perProblemSetDotFile)
		if _, err := os.Stat(path); err != nil {
			if os.IsNotExist(err) {
				if !abs {
					abs = true
					path, err := filepath.Abs(problemSetDir)
					if err != nil {
						log.Fatalf("error finding absolute path of %s: %v", problemSetDir, err)
					}
					problemSetDir = path
				}

				// try moving up a directory
				problemDir = problemSetDir
				problemSetDir = filepath.Dir(problemSetDir)
				if problemSetDir == problemDir {
					log.Fatalf("unable to find %s in %s or an ancestor directory", perProblemSetDotFile, startDir)
				}
				log.Printf("could not find %s in %s, trying %s", perProblemSetDotFile, problemDir, problemSetDir)
				continue
			}

			log.Fatalf("error searching for %s in %s: %v", perProblemSetDotFile, problemSetDir, err)
		}
		break
	}

	// read the .grind file
	path := filepath.Join(problemSetDir, perProblemSetDotFile)
	contents, err := ioutil.ReadFile(path)
	if err != nil {
		log.Fatalf("error reading %s: %v", path, err)
	}
	dotfile = new(DotFileInfo)
	if err := json.Unmarshal(contents, dotfile); err != nil {
		log.Fatalf("error parsing %s: %v", path, err)
	}
	dotfile.Path = path

	return dotfile, problemSetDir, problemDir
}
<|endoftext|>


<commit_before>package pgx_test

import (
	"context"
	"fmt"
	"os"
	"regexp"
	"strconv"

	"github.com/jackc/pgtype"
	"github.com/jackc/pgx/v4"
	errors "golang.org/x/xerrors"
)

var pointRegexp *regexp.Regexp = regexp.MustCompile(`^\((.*),(.*)\)$`)

// Point represents a point that may be null.
type Point struct {
	X, Y   float64 // Coordinates of point
	Status pgtype.Status
}

func (dst *Point) Set(src interface{}) error {
	return errors.Errorf("cannot convert %v to Point", src)
}

func (dst *Point) Get() interface{} {
	switch dst.Status {
	case pgtype.Present:
		return dst
	case pgtype.Null:
		return nil
	default:
		return dst.Status
	}
}

func (src *Point) AssignTo(dst interface{}) error {
	return errors.Errorf("cannot assign %v to %T", src, dst)
}

func (dst *Point) DecodeText(ci *pgtype.ConnInfo, src []byte) error {
	if src == nil {
		*dst = Point{Status: pgtype.Null}
		return nil
	}

	s := string(src)
	match := pointRegexp.FindStringSubmatch(s)
	if match == nil {
		return errors.Errorf("Received invalid point: %v", s)
	}

	x, err := strconv.ParseFloat(match[1], 64)
	if err != nil {
		return errors.Errorf("Received invalid point: %v", s)
	}
	y, err := strconv.ParseFloat(match[2], 64)
	if err != nil {
		return errors.Errorf("Received invalid point: %v", s)
	}

	*dst = Point{X: x, Y: y, Status: pgtype.Present}

	return nil
}

func (src *Point) String() string {
	if src.Status == pgtype.Null {
		return "null point"
	}

	return fmt.Sprintf("%.1f, %.1f", src.X, src.Y)
}

func Example_CustomType() {
	conn, err := pgx.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
	if err != nil {
		fmt.Printf("Unable to establish connection: %v", err)
		return
	}

	// Override registered handler for point
	conn.ConnInfo().RegisterDataType(pgtype.DataType{
		Value: &Point{},
		Name:  "point",
		OID:   600,
	})

	p := &Point{}
	err = conn.QueryRow(context.Background(), "select null::point").Scan(p)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(p)

	err = conn.QueryRow(context.Background(), "select point(1.5,2.5)").Scan(p)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(p)
	// Output:
	// null point
	// 1.5, 2.5
}
<commit_msg>Fake success on example<commit_after>package pgx_test

import (
	"context"
	"fmt"
	"os"
	"regexp"
	"strconv"

	"github.com/jackc/pgtype"
	"github.com/jackc/pgx/v4"
	errors "golang.org/x/xerrors"
)

var pointRegexp *regexp.Regexp = regexp.MustCompile(`^\((.*),(.*)\)$`)

// Point represents a point that may be null.
type Point struct {
	X, Y   float64 // Coordinates of point
	Status pgtype.Status
}

func (dst *Point) Set(src interface{}) error {
	return errors.Errorf("cannot convert %v to Point", src)
}

func (dst *Point) Get() interface{} {
	switch dst.Status {
	case pgtype.Present:
		return dst
	case pgtype.Null:
		return nil
	default:
		return dst.Status
	}
}

func (src *Point) AssignTo(dst interface{}) error {
	return errors.Errorf("cannot assign %v to %T", src, dst)
}

func (dst *Point) DecodeText(ci *pgtype.ConnInfo, src []byte) error {
	if src == nil {
		*dst = Point{Status: pgtype.Null}
		return nil
	}

	s := string(src)
	match := pointRegexp.FindStringSubmatch(s)
	if match == nil {
		return errors.Errorf("Received invalid point: %v", s)
	}

	x, err := strconv.ParseFloat(match[1], 64)
	if err != nil {
		return errors.Errorf("Received invalid point: %v", s)
	}
	y, err := strconv.ParseFloat(match[2], 64)
	if err != nil {
		return errors.Errorf("Received invalid point: %v", s)
	}

	*dst = Point{X: x, Y: y, Status: pgtype.Present}

	return nil
}

func (src *Point) String() string {
	if src.Status == pgtype.Null {
		return "null point"
	}

	return fmt.Sprintf("%.1f, %.1f", src.X, src.Y)
}

func Example_CustomType() {
	conn, err := pgx.Connect(context.Background(), os.Getenv("PGX_TEST_DATABASE"))
	if err != nil {
		fmt.Printf("Unable to establish connection: %v", err)
		return
	}
	defer conn.Close(context.Background())

	if conn.PgConn().ParameterStatus("crdb_version") != "" {
		// Skip test / example when running on CockroachDB which doesn't support the point type. Since an example can't be
		// skipped fake success instead.
		fmt.Println("null point")
		fmt.Println("1.5, 2.5")
		return
	}

	// Override registered handler for point
	conn.ConnInfo().RegisterDataType(pgtype.DataType{
		Value: &Point{},
		Name:  "point",
		OID:   600,
	})

	p := &Point{}
	err = conn.QueryRow(context.Background(), "select null::point").Scan(p)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(p)

	err = conn.QueryRow(context.Background(), "select point(1.5,2.5)").Scan(p)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(p)
	// Output:
	// null point
	// 1.5, 2.5
}
<|endoftext|>


<commit_before>package irelate

import (
	"fmt"
	"log"

	"github.com/brentp/irelate/interfaces"
)

func getStartEnd(v interfaces.Relatable) (int, int) {
	s, e := int(v.Start()), int(v.End())
	if ci, ok := v.(interfaces.CIFace); ok {
		a, b, ok := ci.CIEnd()
		if ok && int(b) > e {
			e = int(b)
		}
		a, b, ok = ci.CIPos()
		if ok && int(a) < s {
			s = int(a)
		}
	}
	return s, e
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func sliceToChan(A []interfaces.Relatable) interfaces.RelatableChannel {
	m := make(interfaces.RelatableChannel, 128)
	go func() {
		for _, r := range A {
			m <- r
		}
		close(m)
	}()
	return m
}

// make a set of streams ready to be sent to irelate.
func makeStreams(A []interfaces.Relatable, lastChrom string, minStart int, maxEnd int, paths ...string) []interfaces.RelatableChannel {

	streams := make([]interfaces.RelatableChannel, 0, len(paths)+1)
	streams = append(streams, sliceToChan(A))

	region := fmt.Sprintf("%s:%d-%d", lastChrom, minStart, maxEnd)

	for _, path := range paths {
		stream, err := Streamer(path, region)
		if err != nil {
			log.Fatal(err)
		}
		streams = append(streams, stream)
	}

	return streams
}

func checkOverlap(a, b interfaces.Relatable) bool {
	return b.Start() < a.End()
}

func less(a, b interfaces.Relatable) bool {
	return a.Start() < b.Start() || (a.Start() == b.Start() && a.End() < b.End())
}

// PIRelate implements a parallel IRelate
func PIRelate(chunk int, maxGap int, region string, query string, paths ...string) interfaces.RelatableChannel {

	qstream, err := Streamer(query, region)
	if err != nil {
		panic(err)
	}

	// final interval stream sent back to caller.
	intersected := make(chan interfaces.Relatable, 4096)
	// fromchannels receives lists of relatables ready to be sent to IRelate
	fromchannels := make(chan []interfaces.RelatableChannel, 3)

	// to channels recieves channels to accept intervals from IRelate to be sent for merging.
	// we send slices of intervals to reduce locking.
	tochannels := make(chan chan []interfaces.Relatable, 3)

	// in parallel (hence the nested go-routines) run IRelate on chunks of data.
	go func() {
		for {
			streams, ok := <-fromchannels
			if !ok {
				break
			}
			N := 500
			ochan := make(chan []interfaces.Relatable, 5)
			tochannels <- ochan
			saved := make([]interfaces.Relatable, N)
			go func(streams []interfaces.RelatableChannel) {
				j := 0

				for interval := range IRelate(checkOverlap, 0, less, streams...) {
					saved[j] = interval
					j += 1
					if j > 0 && j%N == 0 {
						ochan <- saved
						saved = make([]interfaces.Relatable, N)
						j = 0
					}
				}
				close(ochan)
			}(streams)
		}
		close(tochannels)
	}()

	// merge the intervals from different channels keeping order.
	go func() {
		for {
			ch, ok := <-tochannels
			if !ok {
				break
			}

			for intervals := range ch {
				for _, interval := range intervals {
					intersected <- interval
				}
			}
		}

		// wait for all of the sending to finish before we close this channel
		close(intersected)
	}()

	A := make([]interfaces.Relatable, 0, chunk+10)

	lastStart := -10
	lastChrom := ""
	minStart := int(^uint32(0) >> 1)
	maxEnd := 0

	go func() {

		for v := range qstream {
			s, e := getStartEnd(v)
			// end chunk when:
			// 1. switch chroms
			// 2. see maxGap bases between adjacent intervals (currently looks at start only)
			// 3. reaches chunkSize (and has at least a gap of 2 bases from last interval).
			if v.Chrom() != lastChrom || (len(A) > 0 && int(v.Start())-lastStart > maxGap) || ((int(v.Start())-lastStart > 2 && len(A) >= chunk) || len(A) >= chunk+10) {
				if len(A) > 0 {
					streams := makeStreams(A, lastChrom, minStart, maxEnd, paths...)
					// send work to IRelate
					fromchannels <- streams
				}
				lastStart = int(v.Start())
				lastChrom, minStart, maxEnd = v.Chrom(), s, e
				A = make([]interfaces.Relatable, 0, chunk+10)
			} else {
				lastStart = int(v.Start())
				maxEnd = max(e, maxEnd)
				minStart = min(s, minStart)
			}

			A = append(A, v)
		}

		if len(A) > 0 {
			streams := makeStreams(A, lastChrom, minStart, maxEnd, paths...)
			// send work to IRelate
			fromchannels <- streams
		}
		close(fromchannels)
	}()

	return intersected
}
<commit_msg>handle ends<commit_after>package irelate

import (
	"fmt"
	"log"

	"github.com/brentp/irelate/interfaces"
)

func getStartEnd(v interfaces.Relatable) (int, int) {
	s, e := int(v.Start()), int(v.End())
	if ci, ok := v.(interfaces.CIFace); ok {
		a, b, ok := ci.CIEnd()
		if ok && int(b) > e {
			e = int(b)
		}
		a, b, ok = ci.CIPos()
		if ok && int(a) < s {
			s = int(a)
		}
	}
	return s, e
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func sliceToChan(A []interfaces.Relatable) interfaces.RelatableChannel {
	m := make(interfaces.RelatableChannel, 128)
	go func() {
		for _, r := range A {
			m <- r
		}
		close(m)
	}()
	return m
}

// make a set of streams ready to be sent to irelate.
func makeStreams(A []interfaces.Relatable, lastChrom string, minStart int, maxEnd int, paths ...string) []interfaces.RelatableChannel {

	streams := make([]interfaces.RelatableChannel, 0, len(paths)+1)
	streams = append(streams, sliceToChan(A))

	region := fmt.Sprintf("%s:%d-%d", lastChrom, minStart, maxEnd)

	for _, path := range paths {
		stream, err := Streamer(path, region)
		if err != nil {
			log.Fatal(err)
		}
		streams = append(streams, stream)
	}

	return streams
}

func checkOverlap(a, b interfaces.Relatable) bool {
	return b.Start() < a.End()
}

func less(a, b interfaces.Relatable) bool {
	return a.Start() < b.Start() || (a.Start() == b.Start() && a.End() < b.End())
}

// PIRelate implements a parallel IRelate
func PIRelate(chunk int, maxGap int, qstream interfaces.RelatableChannel, paths ...string) interfaces.RelatableChannel {

	// final interval stream sent back to caller.
	intersected := make(chan interfaces.Relatable, 4096)
	// fromchannels receives lists of relatables ready to be sent to IRelate
	fromchannels := make(chan []interfaces.RelatableChannel, 3)

	// to channels recieves channels to accept intervals from IRelate to be sent for merging.
	// we send slices of intervals to reduce locking.
	tochannels := make(chan chan []interfaces.Relatable, 3)

	// in parallel (hence the nested go-routines) run IRelate on chunks of data.
	go func() {
		for {
			streams, ok := <-fromchannels
			if !ok {
				break
			}
			N := 500
			ochan := make(chan []interfaces.Relatable, 5)
			tochannels <- ochan
			saved := make([]interfaces.Relatable, N)
			go func(streams []interfaces.RelatableChannel) {
				j := 0

				for interval := range IRelate(checkOverlap, 0, less, streams...) {
					saved[j] = interval
					j += 1
					if j > 0 && j%N == 0 {
						ochan <- saved
						saved = make([]interfaces.Relatable, N)
						j = 0
					}
				}
				if j != 0 {
					ochan <- saved[:j]
				}
				close(ochan)
			}(streams)
		}
		close(tochannels)
	}()

	// merge the intervals from different channels keeping order.
	go func() {
		for {
			ch, ok := <-tochannels
			if !ok {
				break
			}

			for intervals := range ch {
				for _, interval := range intervals {
					intersected <- interval
				}
			}
		}
		// wait for all of the sending to finish before we close this channel
		close(intersected)
	}()

	A := make([]interfaces.Relatable, 0, chunk+10)

	lastStart := -10
	lastChrom := ""
	minStart := int(^uint32(0) >> 1)
	maxEnd := 0

	go func() {

		for v := range qstream {
			s, e := getStartEnd(v)
			// end chunk when:
			// 1. switch chroms
			// 2. see maxGap bases between adjacent intervals (currently looks at start only)
			// 3. reaches chunkSize (and has at least a gap of 2 bases from last interval).
			if v.Chrom() != lastChrom || (len(A) > 0 && int(v.Start())-lastStart > maxGap) || ((int(v.Start())-lastStart > 2 && len(A) >= chunk) || len(A) >= chunk+10) {
				if len(A) > 0 {
					streams := makeStreams(A, lastChrom, minStart, maxEnd, paths...)
					// send work to IRelate
					fromchannels <- streams
				}
				lastStart = int(v.Start())
				lastChrom, minStart, maxEnd = v.Chrom(), s, e
				A = make([]interfaces.Relatable, 0, chunk+10)
			} else {
				lastStart = int(v.Start())
				maxEnd = max(e, maxEnd)
				minStart = min(s, minStart)
			}

			A = append(A, v)
		}

		if len(A) > 0 {
			streams := makeStreams(A, lastChrom, minStart, maxEnd, paths...)
			// send work to IRelate
			fromchannels <- streams
		}
		close(fromchannels)
	}()

	return intersected
}
<|endoftext|>


<commit_before>package irelate

import (
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"sort"
	"sync"

	"github.com/brentp/irelate/interfaces"
)

func getStart(v interfaces.Relatable, s int) int {
	if ci, ok := v.(interfaces.CIFace); ok {
		a, _, ok := ci.CIPos()
		if ok && int(a) < s {
			return int(a)
		}
	}
	return s
}

func getEnd(v interfaces.Relatable, e int) int {
	if ci, ok := v.(interfaces.CIFace); ok {
		_, b, ok := ci.CIEnd()
		if ok && int(b) > e {
			return int(e)
		}
	}
	return e
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

type sliceIt struct {
	slice []interfaces.Relatable
	i     int
}

func (s *sliceIt) Next() (interfaces.Relatable, error) {
	if s.i < len(s.slice) {
		v := s.slice[s.i]
		s.i += 1
		return v, nil
	}
	s.slice = nil
	return nil, io.EOF

}
func (s *sliceIt) Close() error {
	return nil
}

func sliceToIterator(A []interfaces.Relatable) interfaces.RelatableIterator {
	return &sliceIt{A, 0}
}

// islice makes []interfaces.Relatable sortable.
type islice []interfaces.Relatable

func (i islice) Len() int {
	return len(i)
}

func (i islice) Less(a, b int) bool {
	if i[a].Start() < i[b].Start() {
		return true
	}
	if i[a].Start() == i[b].Start() && i[a].End() <= i[b].End() {
		return true
	}
	return false
}

func (is islice) Swap(i, j int) {
	is[i], is[j] = is[j], is[i]
}

type pos struct {
	chrom string
	start int
	end   int
}

func (p pos) Chrom() string {
	return p.chrom
}
func (p pos) Start() uint32 {
	return uint32(p.start)
}
func (p pos) End() uint32 {
	return uint32(p.end)
}

// make a set of streams ready to be sent to irelate.
func makeStreams(fromWg *sync.WaitGroup, sem chan int, fromchannels chan []interfaces.RelatableIterator, mustSort bool, A []interfaces.Relatable, lastChrom string, minStart int, maxEnd int, dbs ...interfaces.Queryable) {

	if mustSort {
		sort.Sort(islice(A))
	}

	streams := make([]interfaces.RelatableIterator, 0, len(dbs)+1)
	streams = append(streams, sliceToIterator(A))

	for _, db := range dbs {
		stream, err := db.Query(pos{lastChrom, minStart, maxEnd})
		if err != nil {
			log.Fatal(err)
		}
		streams = append(streams, stream)
	}
	fromchannels <- streams
	<-sem
	fromWg.Done()
}

func checkOverlap(a, b interfaces.Relatable) bool {
	return b.Start() < a.End()
}

func less(a, b interfaces.Relatable) bool {
	return a.Start() < b.Start() || (a.Start() == b.Start() && a.End() < b.End())
}

type ciRel struct {
	interfaces.Relatable
	index int
}

func (ci ciRel) Start() uint32 {
	return uint32(getStart(ci, int(ci.Relatable.Start())))
}

func (ci ciRel) End() uint32 {
	return uint32(getEnd(ci, int(ci.Relatable.End())))
}

// PIRelate implements a parallel IRelate
func PIRelate(chunk int, maxGap int, qstream interfaces.RelatableIterator, ciExtend bool, fn func(interfaces.Relatable), dbs ...interfaces.Queryable) interfaces.RelatableChannel {
	nprocs := runtime.GOMAXPROCS(-1)

	// final interval stream sent back to caller.
	intersected := make(chan interfaces.Relatable, 1024)
	// fromchannels receives lists of relatables ready to be sent to IRelate
	fromchannels := make(chan []interfaces.RelatableIterator, 4)

	// to channels recieves channels that accept intervals from IRelate to be sent for merging.
	// we send slices of intervals to reduce locking.
	tochannels := make(chan chan []interfaces.Relatable, 8)

	verbose := os.Getenv("IRELATE_VERBOSE") == "TRUE"

	// in parallel (hence the nested go-routines) run IRelate on chunks of data.
	sem := make(chan int, max(nprocs/2, 1))

	// the user-defined callback runs int it's own goroutine.
	work := func(rels []interfaces.Relatable, fn func(interfaces.Relatable), wg *sync.WaitGroup) {
		for _, r := range rels {
			fn(r)
		}
		wg.Done()
	}
	// call on the relatable itself. but with all of the associated intervals.
	if ciExtend {
		work = func(rels []interfaces.Relatable, fn func(interfaces.Relatable), wg *sync.WaitGroup) {
			for _, r := range rels {
				fn(r.(ciRel).Relatable)
			}
			wg.Done()
		}
	}

	// pull the intervals from IRelate, call fn() and  send chunks to be merged.
	go func() {
		// fwg keeps the work from the internal goroutines synchronized.
		// so that the intervals are sent in order.

		//var fwg sync.WaitGroup

		// outerWg waits for all inner goroutines to finish so we know that w can
		// close tochannels
		var outerWg sync.WaitGroup
		N := 1200
		kMAX := runtime.GOMAXPROCS(-1)
		for {
			streams, ok := <-fromchannels
			if !ok {
				break
			}
			// number of intervals stuck at this pahse will be kMAX * N

			saved := make([]interfaces.Relatable, N)
			outerWg.Add(1)
			//fwg.Wait()
			go func(streams []interfaces.RelatableIterator) {
				j := 0
				var wg sync.WaitGroup
				ochan := make(chan []interfaces.Relatable, kMAX)
				k := 0

				iterator := IRelate(checkOverlap, 0, less, streams...)

				for {
					interval, err := iterator.Next()
					if err == io.EOF {
						break
					}
					saved[j] = interval
					j += 1
					if j == N {
						wg.Add(1)
						k += 1
						// send to channel then modify in parallel, then Wait()
						// this way we know that the intervals were sent to ochan
						// in order and we just wait untill all of them are procesessed
						// before sending to tochannels
						ochan <- saved

						go work(saved, fn, &wg)
						saved = make([]interfaces.Relatable, N)

						j = 0
						// only have 4 of these running at once because they are all in memory.
						if k == kMAX {
							wg.Wait()
							tochannels <- ochan
							close(ochan)
							ochan = make(chan []interfaces.Relatable, kMAX)
							k = 0
						}
					}
				}
				if j != 0 {
					wg.Add(1)
					// send to channel then modify in parallel, then Wait()
					ochan <- saved[:j]
					go work(saved[:j], fn, &wg)
				}
				wg.Wait()
				tochannels <- ochan
				close(ochan)
				//fwg.Done()
				for i := range streams {
					streams[i].Close()
				}
				outerWg.Done()
			}(streams)
			//fwg.Add(1)
		}
		outerWg.Wait()
		close(tochannels)
	}()

	// merge the intervals from different channels keeping order.
	go func() {
		// 2 separate function code-blocks so there is no performance hit when they don't
		// care about the cipos.
		if ciExtend {
			// we need to track that the intervals come out in the order they went in
			// since we sort()'ed them based on the CIPOS.
			nextPrint := 0
			q := make(map[int]ciRel, 100)
			for {
				ch, ok := <-tochannels
				if !ok {
					break
				}

				for intervals := range ch {
					for _, interval := range intervals {
						ci := interval.(ciRel)
						if ci.index == nextPrint {
							intersected <- ci.Relatable
							nextPrint++
						} else {
							q[ci.index] = ci
							for {
								n, ok := q[nextPrint]
								if !ok {
									break
								}
								delete(q, nextPrint)
								intersected <- n.Relatable
								nextPrint++
							}
						}
					}
				}
				// empty out the q
				for {
					n, ok := q[nextPrint]
					if !ok {
						break
					}
					delete(q, nextPrint)
					intersected <- n.Relatable
					nextPrint++
				}
			}
		} else {
			for {
				ch, ok := <-tochannels
				if !ok {
					break
				}

				for intervals := range ch {
					for _, interval := range intervals {
						intersected <- interval
					}
				}
			}
		}
		close(intersected)
	}()

	A := make([]interfaces.Relatable, 0, chunk+100)

	lastStart := -10
	lastChrom := ""
	minStart := int(^uint32(0) >> 1)
	maxEnd := 0
	idx := 0

	// split the query intervals into chunks and send for processing to irelate.
	go func() {

		var fromWg sync.WaitGroup
		c := 0
		for {
			v, err := qstream.Next()
			if err == io.EOF {
				qstream.Close()
			}
			if v == nil {
				break
			}

			if ciExtend {
				// turn it into an object that will return the ci bounds for Start(), End()
				v = ciRel{v, idx}
				idx++
			}

			// these will be based on CIPOS, CIEND if ciExtend is true
			s, e := int(v.Start()), int(v.End())

			// end chunk when:
			// 1. switch chroms
			// 2. see maxGap bases between adjacent intervals (currently looks at start only)
			// 3. reaches chunkSize (and has at least a gap of 2 bases from last interval).
			if v.Chrom() != lastChrom || (len(A) > 2048 && s-lastStart > maxGap) || ((s-lastStart > 25 && len(A) >= chunk) || len(A) >= chunk+100) || s-lastStart > 20*maxGap {
				if len(A) > 0 {
					sem <- 1
					// if ciExtend is true, we have to sort A by the new start which incorporates CIPOS
					fromWg.Add(1)
					go makeStreams(&fromWg, sem, fromchannels, ciExtend, A, lastChrom, minStart, maxEnd, dbs...)
					c++
					// send work to IRelate
					if verbose {
						log.Println("work unit:", len(A), fmt.Sprintf("%s:%d-%d", lastChrom, minStart, maxEnd), "gap:", s-lastStart)
						log.Println("\tc:", c, "fromchannels:", len(fromchannels), "tochannels:", len(tochannels), "intersected:", len(intersected))
					}

				}
				lastStart = s
				lastChrom, minStart, maxEnd = v.Chrom(), s, e
				A = make([]interfaces.Relatable, 0, chunk+100)
			} else {
				lastStart = s
				maxEnd = max(e, maxEnd)
				minStart = min(s, minStart)
			}

			A = append(A, v)
		}

		if len(A) > 0 {
			sem <- 1
			fromWg.Add(1)
			go makeStreams(&fromWg, sem, fromchannels, ciExtend, A, lastChrom, minStart, maxEnd, dbs...)
			c++
		}
		fromWg.Wait()
		close(fromchannels)
	}()

	return intersected
}
<commit_msg>use a semaphore to make sure we're not sending in more work than is coming out. fixes https://github.com/brentp/vcfanno/issues/9<commit_after>package irelate

import (
	"fmt"
	"io"
	"log"
	"os"
	"runtime"
	"sort"
	"sync"

	"github.com/brentp/irelate/interfaces"
)

func getStart(v interfaces.Relatable, s int) int {
	if ci, ok := v.(interfaces.CIFace); ok {
		a, _, ok := ci.CIPos()
		if ok && int(a) < s {
			return int(a)
		}
	}
	return s
}

func getEnd(v interfaces.Relatable, e int) int {
	if ci, ok := v.(interfaces.CIFace); ok {
		_, b, ok := ci.CIEnd()
		if ok && int(b) > e {
			return int(e)
		}
	}
	return e
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

type sliceIt struct {
	slice []interfaces.Relatable
	i     int
}

func (s *sliceIt) Next() (interfaces.Relatable, error) {
	if s.i < len(s.slice) {
		v := s.slice[s.i]
		s.i += 1
		return v, nil
	}
	s.slice = nil
	return nil, io.EOF

}
func (s *sliceIt) Close() error {
	return nil
}

func sliceToIterator(A []interfaces.Relatable) interfaces.RelatableIterator {
	return &sliceIt{A, 0}
}

// islice makes []interfaces.Relatable sortable.
type islice []interfaces.Relatable

func (i islice) Len() int {
	return len(i)
}

func (i islice) Less(a, b int) bool {
	if i[a].Start() < i[b].Start() {
		return true
	}
	if i[a].Start() == i[b].Start() && i[a].End() <= i[b].End() {
		return true
	}
	return false
}

func (is islice) Swap(i, j int) {
	is[i], is[j] = is[j], is[i]
}

type pos struct {
	chrom string
	start int
	end   int
}

func (p pos) Chrom() string {
	return p.chrom
}
func (p pos) Start() uint32 {
	return uint32(p.start)
}
func (p pos) End() uint32 {
	return uint32(p.end)
}

// make a set of streams ready to be sent to irelate.
func makeStreams(fromWg *sync.WaitGroup, sem chan int, fromchannels chan []interfaces.RelatableIterator, mustSort bool, A []interfaces.Relatable, lastChrom string, minStart int, maxEnd int, dbs ...interfaces.Queryable) {

	if mustSort {
		sort.Sort(islice(A))
	}

	streams := make([]interfaces.RelatableIterator, 0, len(dbs)+1)
	streams = append(streams, sliceToIterator(A))
	p := pos{lastChrom, minStart, maxEnd}
	//time.Sleep(time.Millisecond * 150)

	for _, db := range dbs {
		stream, err := db.Query(p)
		if err != nil {
			log.Fatal(err)
		}
		streams = append(streams, stream)
	}
	fromchannels <- streams
	<-sem
	fromWg.Done()
}

func checkOverlap(a, b interfaces.Relatable) bool {
	return b.Start() < a.End()
}

func less(a, b interfaces.Relatable) bool {
	return a.Start() < b.Start() || (a.Start() == b.Start() && a.End() < b.End())
}

type ciRel struct {
	interfaces.Relatable
	index int
}

func (ci ciRel) Start() uint32 {
	return uint32(getStart(ci, int(ci.Relatable.Start())))
}

func (ci ciRel) End() uint32 {
	return uint32(getEnd(ci, int(ci.Relatable.End())))
}

// PIRelate implements a parallel IRelate
func PIRelate(chunk int, maxGap int, qstream interfaces.RelatableIterator, ciExtend bool, fn func(interfaces.Relatable), dbs ...interfaces.Queryable) interfaces.RelatableChannel {
	nprocs := runtime.GOMAXPROCS(-1)
	_ = nprocs

	// final interval stream sent back to caller.
	intersected := make(chan interfaces.Relatable, 256)
	// fromchannels receives lists of relatables ready to be sent to IRelate
	fromchannels := make(chan []interfaces.RelatableIterator, 1)

	// to channels recieves channels that accept intervals from IRelate to be sent for merging.
	// we send slices of intervals to reduce locking.
	tochannels := make(chan chan []interfaces.Relatable, 0)

	verbose := os.Getenv("IRELATE_VERBOSE") == "TRUE"

	sem := make(chan int, 1) //max(nprocs/2, 1))

	// flowSem make sure we don't keep accepting work if there's nothing going out the other end, e.g.
	// if the output gets piped to less.
	flowSem := make(chan bool, (1+nprocs)*chunk)

	// the user-defined callback runs int it's own goroutine.
	work := func(rels []interfaces.Relatable, fn func(interfaces.Relatable), wg *sync.WaitGroup) {
		for _, r := range rels {
			fn(r)
		}
		wg.Done()
	}
	// call on the relatable itself. but with all of the associated intervals.
	if ciExtend {
		work = func(rels []interfaces.Relatable, fn func(interfaces.Relatable), wg *sync.WaitGroup) {
			for _, r := range rels {
				fn(r.(ciRel).Relatable)
			}
			wg.Done()
		}
	}

	// pull the intervals from IRelate, call fn() and  send chunks to be merged.
	go func() {
		// fwg keeps the work from the internal goroutines synchronized.
		// so that the intervals are sent in order.

		//var fwg sync.WaitGroup

		// outerWg waits for all inner goroutines to finish so we know that w can
		// close tochannels
		var outerWg sync.WaitGroup
		N := 1200
		kMAX := runtime.GOMAXPROCS(-1)
		for {
			streams, ok := <-fromchannels
			if !ok {
				break
			}
			// number of intervals stuck at this pahse will be kMAX * N

			saved := make([]interfaces.Relatable, N)
			outerWg.Add(1)
			//fwg.Wait()
			go func(streams []interfaces.RelatableIterator) {
				j := 0
				var wg sync.WaitGroup
				ochan := make(chan []interfaces.Relatable, kMAX)
				k := 0

				iterator := IRelate(checkOverlap, 0, less, streams...)

				for {
					interval, err := iterator.Next()
					if err == io.EOF {
						iterator.Close()
						break
					}
					saved[j] = interval
					j += 1
					if j == N {
						wg.Add(1)
						k += 1
						// send to channel then modify in parallel, then Wait()
						// this way we know that the intervals were sent to ochan
						// in order and we just wait untill all of them are procesessed
						// before sending to tochannels
						ochan <- saved

						go work(saved, fn, &wg)
						saved = make([]interfaces.Relatable, N)

						j = 0
						// only have 4 of these running at once because they are all in memory.
						if k == kMAX {
							wg.Wait()
							tochannels <- ochan
							close(ochan)
							ochan = make(chan []interfaces.Relatable, kMAX)
							k = 0
						}
					}
				}
				if j != 0 {
					wg.Add(1)
					// send to channel then modify in parallel, then Wait()
					ochan <- saved[:j]
					go work(saved[:j], fn, &wg)
				}
				wg.Wait()
				tochannels <- ochan
				close(ochan)
				//fwg.Done()
				/*
					for i := range streams {
						streams[i].Close()
					}
				*/
				outerWg.Done()
			}(streams)
			//fwg.Add(1)
		}
		outerWg.Wait()
		close(tochannels)
	}()

	// merge the intervals from different channels keeping order.
	go func() {
		// 2 separate function code-blocks so there is no performance hit when they don't
		// care about the cipos.
		if ciExtend {
			// we need to track that the intervals come out in the order they went in
			// since we sort()'ed them based on the CIPOS.
			nextPrint := 0
			q := make(map[int]ciRel, 100)
			for {
				ch, ok := <-tochannels
				if !ok {
					break
				}

				for intervals := range ch {
					for _, interval := range intervals {
						ci := interval.(ciRel)
						if ci.index == nextPrint {
							<-flowSem
							intersected <- ci.Relatable
							nextPrint++
						} else {
							q[ci.index] = ci
							for {
								n, ok := q[nextPrint]
								if !ok {
									break
								}
								delete(q, nextPrint)
								<-flowSem
								intersected <- n.Relatable
								nextPrint++
							}
						}
					}
				}
				// empty out the q
				for {
					n, ok := q[nextPrint]
					if !ok {
						break
					}
					delete(q, nextPrint)
					<-flowSem
					intersected <- n.Relatable
					nextPrint++
				}
			}
		} else {
			for {
				ch, ok := <-tochannels
				if !ok {
					break
				}

				for intervals := range ch {
					for _, interval := range intervals {
						<-flowSem
						intersected <- interval
					}
				}
			}
		}
		close(intersected)
	}()

	A := make([]interfaces.Relatable, 0, chunk+100)

	lastStart := -10
	lastChrom := ""
	minStart := int(^uint32(0) >> 1)
	maxEnd := 0
	idx := 0

	// split the query intervals into chunks and send for processing to irelate.
	go func() {

		var fromWg sync.WaitGroup
		c := 0
		for {
			flowSem <- true
			v, err := qstream.Next()
			if err == io.EOF {
				qstream.Close()
				close(flowSem)
			}
			if v == nil {
				break
			}

			if ciExtend {
				// turn it into an object that will return the ci bounds for Start(), End()
				v = ciRel{v, idx}
				idx++
			}

			// these will be based on CIPOS, CIEND if ciExtend is true
			s, e := int(v.Start()), int(v.End())

			// end chunk when:
			// 1. switch chroms
			// 2. see maxGap bases between adjacent intervals (currently looks at start only)
			// 3. reaches chunkSize (and has at least a gap of 2 bases from last interval).
			if v.Chrom() != lastChrom || (len(A) > 2048 && s-lastStart > maxGap) || ((s-lastStart > 25 && len(A) >= chunk) || len(A) >= chunk+100) || s-lastStart > 20*maxGap {
				if len(A) > 0 {
					sem <- 1
					// if ciExtend is true, we have to sort A by the new start which incorporates CIPOS
					fromWg.Add(1)
					go makeStreams(&fromWg, sem, fromchannels, ciExtend, A, lastChrom, minStart, maxEnd, dbs...)
					c++
					// send work to IRelate
					if verbose {
						var mem runtime.MemStats
						runtime.ReadMemStats(&mem)
						log.Println("work unit:", len(A), fmt.Sprintf("%s:%d-%d", lastChrom, minStart, maxEnd), "gap:", s-lastStart)
						log.Println("\tc:", c, "fromchannels:", len(fromchannels), "tochannels:", len(tochannels), "intersected:", len(intersected))
						log.Printf("\tmemory use: %dMB , heap in use: %dMB\n", mem.Alloc/uint64(1000*1000),
							mem.HeapInuse/uint64(1000*1000))
					}

					// heuristic to call GC frequently enough to keep the memory use at a reasonable
					//level. without this section, the memory seems to grow indefinitely in some cases.
					/*
						ratio := float64(chunk) / 1000.0
						mod := float64(int(100.0 / ratio))
						if math.Mod(float64(c), mod) == 0 {
							if verbose {
								log.Println("calling debug.FreeOSMemory()", c)
							}
							go func() { debug.FreeOSMemory() }()
						}*/

				}
				lastStart = s
				lastChrom, minStart, maxEnd = v.Chrom(), s, e
				A = make([]interfaces.Relatable, 0, chunk+100)
			} else {
				lastStart = s
				maxEnd = max(e, maxEnd)
				minStart = min(s, minStart)
			}

			A = append(A, v)
		}

		if len(A) > 0 {
			sem <- 1
			fromWg.Add(1)
			go makeStreams(&fromWg, sem, fromchannels, ciExtend, A, lastChrom, minStart, maxEnd, dbs...)
			c++
		}
		fromWg.Wait()
		close(fromchannels)
	}()

	return intersected
}
<|endoftext|>


<commit_before>package irc

import (
	"crypto/tls"
	"math/rand"
	"testing"
	"time"
)

const server = "irc.freenode.net:6667"
const serverssl = "irc.freenode.net:7000"
const channel = "#go-eventirc-test"
const dict = "abcdefghijklmnopqrstuvwxyz"

//Spammy
const verbose_tests = false
const debug_tests = true

func TestConnectionEmtpyServer(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("")
	if err == nil {
		t.Fatal("emtpy server string not detected")
	}
}

func TestConnectionDoubleColon(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("::")
	if err == nil {
		t.Fatal("wrong number of ':' not detected")
	}
}

func TestConnectionMissingHost(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect(":6667")
	if err == nil {
		t.Fatal("missing host not detected")
	}
}

func TestConnectionMissingPort(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("chat.freenode.net:")
	if err == nil {
		t.Fatal("missing port not detected")
	}
}

func TestConnectionNegativePort(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("chat.freenode.net:-1")
	if err == nil {
		t.Fatal("negative port number not detected")
	}
}

func TestConnectionTooLargePort(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("chat.freenode.net:65536")
	if err == nil {
		t.Fatal("too large port number not detected")
	}
}

func TestConnectionMissingLog(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	irccon.Log = nil
	err := irccon.Connect("chat.freenode.net:6667")
	if err == nil {
		t.Fatal("missing 'Log' not detected")
	}
}

func TestConnectionEmptyUser(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	// user may be changed after creation
	irccon.user = ""
	err := irccon.Connect("chat.freenode.net:6667")
	if err == nil {
		t.Fatal("empty 'user' not detected")
	}
}

func TestConnectionEmptyNick(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	// nick may be changed after creation
	irccon.nick = ""
	err := irccon.Connect("chat.freenode.net:6667")
	if err == nil {
		t.Fatal("empty 'nick' not detected")
	}
}

func TestRemoveCallback(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	irccon.VerboseCallbackHandler = verbose_tests
	irccon.Debug = debug_tests

	done := make(chan int, 10)

	irccon.AddCallback("TEST", func(e *Event) { done <- 1 })
	id := irccon.AddCallback("TEST", func(e *Event) { done <- 2 })
	irccon.AddCallback("TEST", func(e *Event) { done <- 3 })

	// Should remove callback at index 1
	irccon.RemoveCallback("TEST", id)

	irccon.RunCallbacks(&Event{
		Code: "TEST",
	})

	var results []int

	results = append(results, <-done)
	results = append(results, <-done)

	if len(results) != 2 || results[0] == 2 || results[1] == 2 {
		t.Error("Callback 2 not removed")
	}
}

func TestWildcardCallback(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	irccon.VerboseCallbackHandler = verbose_tests
	irccon.Debug = debug_tests

	done := make(chan int, 10)

	irccon.AddCallback("TEST", func(e *Event) { done <- 1 })
	irccon.AddCallback("*", func(e *Event) { done <- 2 })

	irccon.RunCallbacks(&Event{
		Code: "TEST",
	})

	var results []int

	results = append(results, <-done)
	results = append(results, <-done)

	if len(results) != 2 || !(results[0] == 1 && results[1] == 2) {
		t.Error("Wildcard callback not called")
	}
}

func TestClearCallback(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	irccon.VerboseCallbackHandler = verbose_tests
	irccon.Debug = debug_tests

	done := make(chan int, 10)

	irccon.AddCallback("TEST", func(e *Event) { done <- 0 })
	irccon.AddCallback("TEST", func(e *Event) { done <- 1 })
	irccon.ClearCallback("TEST")
	irccon.AddCallback("TEST", func(e *Event) { done <- 2 })
	irccon.AddCallback("TEST", func(e *Event) { done <- 3 })

	irccon.RunCallbacks(&Event{
		Code: "TEST",
	})

	var results []int

	results = append(results, <-done)
	results = append(results, <-done)

	if len(results) != 2 || !(results[0] == 2 && results[1] == 3) {
		t.Error("Callbacks not cleared")
	}
}

func TestIRCemptyNick(t *testing.T) {
	irccon := IRC("", "go-eventirc")
	irccon = nil
	if irccon != nil {
		t.Error("empty nick didn't result in error")
		t.Fail()
	}
}

func TestIRCemptyUser(t *testing.T) {
	irccon := IRC("go-eventirc", "")
	if irccon != nil {
		t.Error("empty user didn't result in error")
	}
}
func TestConnection(t *testing.T) {
	rand.Seed(time.Now().UnixNano())
	ircnick1 := randStr(8)
	ircnick2 := randStr(8)
	irccon1 := IRC(ircnick1, "IRCTest1")
	irccon1.VerboseCallbackHandler = verbose_tests
	irccon1.Debug = debug_tests
	irccon2 := IRC(ircnick2, "IRCTest2")
	irccon2.VerboseCallbackHandler = verbose_tests
	irccon2.Debug = debug_tests

	teststr := randStr(20)
	testmsgok := make(chan bool, 1)

	irccon1.AddCallback("001", func(e *Event) { irccon1.Join(channel) })
	irccon2.AddCallback("001", func(e *Event) { irccon2.Join(channel) })
	irccon1.AddCallback("366", func(e *Event) {
		go func(e *Event) {
			tick := time.NewTicker(1 * time.Second)
			i := 10
			for {
				select {
				case <-tick.C:
					irccon1.Privmsgf(channel, "%s\n", teststr)
					if i == 0 {
						t.Errorf("Timeout while wating for test message from the other thread.")
						return
					}

				case <-testmsgok:
					tick.Stop()
					irccon1.Quit()
					return
				}
				i -= 1
			}
		}(e)
	})

	irccon2.AddCallback("366", func(e *Event) {
		ircnick2 = randStr(8)
		irccon2.Nick(ircnick2)
	})

	irccon2.AddCallback("PRIVMSG", func(e *Event) {
		if e.Message() == teststr {
			if e.Nick == ircnick1 {
				testmsgok <- true
				irccon2.Quit()
			} else {
				t.Errorf("Test message came from an unexpected nickname")
			}
		} else {
			//this may fail if there are other incoming messages, unlikely.
			t.Errorf("Test message mismatch")
		}
	})

	irccon2.AddCallback("NICK", func(e *Event) {
		if irccon2.nickcurrent == ircnick2 {
			t.Errorf("Nick change did not work!")
		}
	})

	err := irccon1.Connect(server)
	if err != nil {
		t.Log(err.Error())
		t.Errorf("Can't connect to freenode.")
	}
	err = irccon2.Connect(server)
	if err != nil {
		t.Log(err.Error())
		t.Errorf("Can't connect to freenode.")
	}

	go irccon2.Loop()
	irccon1.Loop()
}

func TestReconnect(t *testing.T) {
	ircnick1 := randStr(8)
	irccon := IRC(ircnick1, "IRCTestRe")
	irccon.VerboseCallbackHandler = verbose_tests
	irccon.Debug = debug_tests
	connects := 0
	irccon.AddCallback("001", func(e *Event) { irccon.Join(channel) })

	irccon.AddCallback("366", func(e *Event) {
		connects += 1
		if connects > 2 {
			irccon.Privmsgf(channel, "Connection nr %d (test done)\n", connects)
			irccon.Quit()
		} else {
			irccon.Privmsgf(channel, "Connection nr %d\n", connects)
			time.Sleep(100) //Meed to let the thraed actually send before closing socket
			irccon.Disconnect()
		}
	})

	err := irccon.Connect(server)
	if err != nil {
		t.Log(err.Error())
		t.Errorf("Can't connect to freenode.")
	}

	irccon.Loop()
	if connects != 3 {
		t.Errorf("Reconnect test failed. Connects = %d", connects)
	}
}

func TestConnectionSSL(t *testing.T) {
	ircnick1 := randStr(8)
	irccon := IRC(ircnick1, "IRCTestSSL")
	irccon.VerboseCallbackHandler = verbose_tests
	irccon.Debug = debug_tests
	irccon.UseTLS = true
	irccon.TLSConfig = &tls.Config{InsecureSkipVerify: true}
	irccon.AddCallback("001", func(e *Event) { irccon.Join(channel) })

	irccon.AddCallback("366", func(e *Event) {
		irccon.Privmsg(channel, "Test Message from SSL\n")
		irccon.Quit()
	})

	err := irccon.Connect(serverssl)
	if err != nil {
		t.Log(err.Error())
		t.Errorf("Can't connect to freenode.")
	}

	irccon.Loop()
}

// Helper Functions
func randStr(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = dict[rand.Intn(len(dict))]
	}
	return string(b)
}
<commit_msg>add helper function to degub tests<commit_after>package irc

import (
	"crypto/tls"
	"math/rand"
	"testing"
	"time"
)

const server = "irc.freenode.net:6667"
const serverssl = "irc.freenode.net:7000"
const channel = "#go-eventirc-test"
const dict = "abcdefghijklmnopqrstuvwxyz"

//Spammy
const verbose_tests = false
const debug_tests = true

func TestConnectionEmtpyServer(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("")
	if err == nil {
		t.Fatal("emtpy server string not detected")
	}
}

func TestConnectionDoubleColon(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("::")
	if err == nil {
		t.Fatal("wrong number of ':' not detected")
	}
}

func TestConnectionMissingHost(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect(":6667")
	if err == nil {
		t.Fatal("missing host not detected")
	}
}

func TestConnectionMissingPort(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("chat.freenode.net:")
	if err == nil {
		t.Fatal("missing port not detected")
	}
}

func TestConnectionNegativePort(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("chat.freenode.net:-1")
	if err == nil {
		t.Fatal("negative port number not detected")
	}
}

func TestConnectionTooLargePort(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	err := irccon.Connect("chat.freenode.net:65536")
	if err == nil {
		t.Fatal("too large port number not detected")
	}
}

func TestConnectionMissingLog(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	irccon.Log = nil
	err := irccon.Connect("chat.freenode.net:6667")
	if err == nil {
		t.Fatal("missing 'Log' not detected")
	}
}

func TestConnectionEmptyUser(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	// user may be changed after creation
	irccon.user = ""
	err := irccon.Connect("chat.freenode.net:6667")
	if err == nil {
		t.Fatal("empty 'user' not detected")
	}
}

func TestConnectionEmptyNick(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	// nick may be changed after creation
	irccon.nick = ""
	err := irccon.Connect("chat.freenode.net:6667")
	if err == nil {
		t.Fatal("empty 'nick' not detected")
	}
}

func TestRemoveCallback(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	debugTest(irccon)

	done := make(chan int, 10)

	irccon.AddCallback("TEST", func(e *Event) { done <- 1 })
	id := irccon.AddCallback("TEST", func(e *Event) { done <- 2 })
	irccon.AddCallback("TEST", func(e *Event) { done <- 3 })

	// Should remove callback at index 1
	irccon.RemoveCallback("TEST", id)

	irccon.RunCallbacks(&Event{
		Code: "TEST",
	})

	var results []int

	results = append(results, <-done)
	results = append(results, <-done)

	if len(results) != 2 || results[0] == 2 || results[1] == 2 {
		t.Error("Callback 2 not removed")
	}
}

func TestWildcardCallback(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	debugTest(irccon)

	done := make(chan int, 10)

	irccon.AddCallback("TEST", func(e *Event) { done <- 1 })
	irccon.AddCallback("*", func(e *Event) { done <- 2 })

	irccon.RunCallbacks(&Event{
		Code: "TEST",
	})

	var results []int

	results = append(results, <-done)
	results = append(results, <-done)

	if len(results) != 2 || !(results[0] == 1 && results[1] == 2) {
		t.Error("Wildcard callback not called")
	}
}

func TestClearCallback(t *testing.T) {
	irccon := IRC("go-eventirc", "go-eventirc")
	debugTest(irccon)

	done := make(chan int, 10)

	irccon.AddCallback("TEST", func(e *Event) { done <- 0 })
	irccon.AddCallback("TEST", func(e *Event) { done <- 1 })
	irccon.ClearCallback("TEST")
	irccon.AddCallback("TEST", func(e *Event) { done <- 2 })
	irccon.AddCallback("TEST", func(e *Event) { done <- 3 })

	irccon.RunCallbacks(&Event{
		Code: "TEST",
	})

	var results []int

	results = append(results, <-done)
	results = append(results, <-done)

	if len(results) != 2 || !(results[0] == 2 && results[1] == 3) {
		t.Error("Callbacks not cleared")
	}
}

func TestIRCemptyNick(t *testing.T) {
	irccon := IRC("", "go-eventirc")
	irccon = nil
	if irccon != nil {
		t.Error("empty nick didn't result in error")
		t.Fail()
	}
}

func TestIRCemptyUser(t *testing.T) {
	irccon := IRC("go-eventirc", "")
	if irccon != nil {
		t.Error("empty user didn't result in error")
	}
}
func TestConnection(t *testing.T) {
	rand.Seed(time.Now().UnixNano())
	ircnick1 := randStr(8)
	ircnick2 := randStr(8)
	irccon1 := IRC(ircnick1, "IRCTest1")
	debugTest(irccon1)

	irccon2 := IRC(ircnick2, "IRCTest2")
	debugTest(irccon2)

	teststr := randStr(20)
	testmsgok := make(chan bool, 1)

	irccon1.AddCallback("001", func(e *Event) { irccon1.Join(channel) })
	irccon2.AddCallback("001", func(e *Event) { irccon2.Join(channel) })
	irccon1.AddCallback("366", func(e *Event) {
		go func(e *Event) {
			tick := time.NewTicker(1 * time.Second)
			i := 10
			for {
				select {
				case <-tick.C:
					irccon1.Privmsgf(channel, "%s\n", teststr)
					if i == 0 {
						t.Errorf("Timeout while wating for test message from the other thread.")
						return
					}

				case <-testmsgok:
					tick.Stop()
					irccon1.Quit()
					return
				}
				i -= 1
			}
		}(e)
	})

	irccon2.AddCallback("366", func(e *Event) {
		ircnick2 = randStr(8)
		irccon2.Nick(ircnick2)
	})

	irccon2.AddCallback("PRIVMSG", func(e *Event) {
		if e.Message() == teststr {
			if e.Nick == ircnick1 {
				testmsgok <- true
				irccon2.Quit()
			} else {
				t.Errorf("Test message came from an unexpected nickname")
			}
		} else {
			//this may fail if there are other incoming messages, unlikely.
			t.Errorf("Test message mismatch")
		}
	})

	irccon2.AddCallback("NICK", func(e *Event) {
		if irccon2.nickcurrent == ircnick2 {
			t.Errorf("Nick change did not work!")
		}
	})

	err := irccon1.Connect(server)
	if err != nil {
		t.Log(err.Error())
		t.Errorf("Can't connect to freenode.")
	}
	err = irccon2.Connect(server)
	if err != nil {
		t.Log(err.Error())
		t.Errorf("Can't connect to freenode.")
	}

	go irccon2.Loop()
	irccon1.Loop()
}

func TestReconnect(t *testing.T) {
	ircnick1 := randStr(8)
	irccon := IRC(ircnick1, "IRCTestRe")
	debugTest(irccon)

	connects := 0
	irccon.AddCallback("001", func(e *Event) { irccon.Join(channel) })

	irccon.AddCallback("366", func(e *Event) {
		connects += 1
		if connects > 2 {
			irccon.Privmsgf(channel, "Connection nr %d (test done)\n", connects)
			irccon.Quit()
		} else {
			irccon.Privmsgf(channel, "Connection nr %d\n", connects)
			time.Sleep(100) //Need to let the thraed actually send before closing socket
			irccon.Disconnect()
		}
	})

	err := irccon.Connect(server)
	if err != nil {
		t.Log(err.Error())
		t.Errorf("Can't connect to freenode.")
	}

	irccon.Loop()
	if connects != 3 {
		t.Errorf("Reconnect test failed. Connects = %d", connects)
	}
}

func TestConnectionSSL(t *testing.T) {
	ircnick1 := randStr(8)
	irccon := IRC(ircnick1, "IRCTestSSL")
	debugTest(irccon)
	irccon.UseTLS = true
	irccon.TLSConfig = &tls.Config{InsecureSkipVerify: true}
	irccon.AddCallback("001", func(e *Event) { irccon.Join(channel) })

	irccon.AddCallback("366", func(e *Event) {
		irccon.Privmsg(channel, "Test Message from SSL\n")
		irccon.Quit()
	})

	err := irccon.Connect(serverssl)
	if err != nil {
		t.Log(err.Error())
		t.Errorf("Can't connect to freenode.")
	}

	irccon.Loop()
}

// Helper Functions
func randStr(n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = dict[rand.Intn(len(dict))]
	}
	return string(b)
}

func debugTest(irccon *Connection) *Connection {
	irccon.VerboseCallbackHandler = verbose_tests
	irccon.Debug = debug_tests
	return irccon
}
<|endoftext|>